From 512910c83f9dd928e2465682a35991f1b2f6aecf Mon Sep 17 00:00:00 2001
From: Matthew Flatt <mflatt@racket-lang.org>
Date: Sat, 28 Oct 2017 10:07:23 -0600
Subject: [PATCH] expander: fix 'module-begin expansion with definition
 contexts

Using `(local-expand <expr> 'module-begin <stops> def-ctx)` didn't
work right, because definitions added to `def-ctx` were not visible.
(While adding definitions before `module-begin` expansion is an unusual
thing to do, there's no reason that it has to fail.)
---
 pkgs/racket-test-core/tests/racket/macro.rktl | 18 ++++++++++++++++++
 racket/src/racket/src/compenv.c               |  5 +++++
 racket/src/racket/src/module.c                |  1 +
 racket/src/racket/src/mzmark_compenv.inc      |  2 ++
 racket/src/racket/src/mzmarksrc.c             |  1 +
 racket/src/racket/src/schpriv.h               |  1 +
 6 files changed, 28 insertions(+)

diff --git a/pkgs/racket-test-core/tests/racket/macro.rktl b/pkgs/racket-test-core/tests/racket/macro.rktl
index 6496813cb2..92ed5487b1 100644
--- a/pkgs/racket-test-core/tests/racket/macro.rktl
+++ b/pkgs/racket-test-core/tests/racket/macro.rktl
@@ -1716,6 +1716,24 @@
                 #'([(b (b ...)) ...] ...)))
              (lambda (exn) (regexp-match? #rx"incompatible ellipsis" (exn-message exn))))
 
+;; ----------------------------------------
+;; Check that expansion to `#%module-begin` is prepared to handle
+;; definition contexts
+
+(module make-definition-context-during-module-begin racket/base
+  (require (for-syntax racket/base))
+  
+  (define-syntax (bind-and-expand stx)
+    (syntax-case stx ()
+      [(_ x form ...)
+       (let ()
+         (define ctx (syntax-local-make-definition-context))
+         (syntax-local-bind-syntaxes (list #'x) #'(λ (stx) (println stx) #'(void)) ctx)
+         (local-expand #'(#%plain-module-begin form ...) 'module-begin '() ctx))]))
+  
+  (module* a #f (bind-and-expand x x))
+  (module* b #f (bind-and-expand x (list x))))
+
 ;; ----------------------------------------
 
 (report-errs)
diff --git a/racket/src/racket/src/compenv.c b/racket/src/racket/src/compenv.c
index 53c6dc3303..addd8bce29 100644
--- a/racket/src/racket/src/compenv.c
+++ b/racket/src/racket/src/compenv.c
@@ -1364,6 +1364,11 @@ scheme_compile_lookup(Scheme_Object *find_id, Scheme_Comp_Env *env, int flags,
 
       if (!frame->vals)
         p += frame->num_bindings;
+
+      if (!frame->next->next && frame->next->intdef_next) {
+        frame = frame->next->intdef_next;
+        continue;
+      }
     }
     
     if (!(flags & SCHEME_OUT_OF_CONTEXT_OK)) {
diff --git a/racket/src/racket/src/module.c b/racket/src/racket/src/module.c
index 17e094660f..263f2dda1a 100644
--- a/racket/src/racket/src/module.c
+++ b/racket/src/racket/src/module.c
@@ -9787,6 +9787,7 @@ static Scheme_Object *do_module_begin_at_phase(Scheme_Object *form, Scheme_Comp_
     cenv = scheme_new_comp_env(env->genv, env->insp, frame_scopes,
                                SCHEME_TOPLEVEL_FRAME | SCHEME_KEEP_SCOPES_FRAME);
     cenv->observer = env->observer;
+    cenv->intdef_next = xenv;
   }
 
   lift_data = scheme_make_vector(3, NULL);
diff --git a/racket/src/racket/src/mzmark_compenv.inc b/racket/src/racket/src/mzmark_compenv.inc
index 1fb9641673..0519d6320a 100644
--- a/racket/src/racket/src/mzmark_compenv.inc
+++ b/racket/src/racket/src/mzmark_compenv.inc
@@ -17,6 +17,7 @@ static int mark_comp_env_MARK(void *p, struct NewGC *gc) {
   gcMARK2(e->prefix, gc);
   gcMARK2(e->next, gc);
   gcMARK2(e->use_scopes_next, gc);
+  gcMARK2(e->intdef_next, gc);
   gcMARK2(e->scopes, gc);
   gcMARK2(e->value_name, gc);
   gcMARK2(e->observer, gc);
@@ -56,6 +57,7 @@ static int mark_comp_env_FIXUP(void *p, struct NewGC *gc) {
   gcFIXUP2(e->prefix, gc);
   gcFIXUP2(e->next, gc);
   gcFIXUP2(e->use_scopes_next, gc);
+  gcFIXUP2(e->intdef_next, gc);
   gcFIXUP2(e->scopes, gc);
   gcFIXUP2(e->value_name, gc);
   gcFIXUP2(e->observer, gc);
diff --git a/racket/src/racket/src/mzmarksrc.c b/racket/src/racket/src/mzmarksrc.c
index f483532c10..180e9c8769 100644
--- a/racket/src/racket/src/mzmarksrc.c
+++ b/racket/src/racket/src/mzmarksrc.c
@@ -1294,6 +1294,7 @@ mark_comp_env {
   gcMARK2(e->prefix, gc);
   gcMARK2(e->next, gc);
   gcMARK2(e->use_scopes_next, gc);
+  gcMARK2(e->intdef_next, gc);
   gcMARK2(e->scopes, gc);
   gcMARK2(e->value_name, gc);
   gcMARK2(e->observer, gc);
diff --git a/racket/src/racket/src/schpriv.h b/racket/src/racket/src/schpriv.h
index 485f9bb49e..580a742fea 100644
--- a/racket/src/racket/src/schpriv.h
+++ b/racket/src/racket/src/schpriv.h
@@ -2878,6 +2878,7 @@ typedef struct Scheme_Comp_Env
 
   struct Scheme_Comp_Env *next;
   struct Scheme_Comp_Env *use_scopes_next; /* fast-forward for use-site scope revert */
+  struct Scheme_Comp_Env *intdef_next; /* when `next` = NULL, can be non-null to continue binding search */
 } Scheme_Comp_Env;
 
 #define LAMBDA_HAS_REST 1