expander: clean up post-expansion scope representation

Instead of three different fields in the context to keep in sync,
encode the possibilities witin a single field.

racket/src/expander/compile/namespace-scope.rkt

@@ -36,7 +36,7 @@
 
 (define (extract-namespace-scopes/values ns)
   (define root-ctx (namespace-get-root-expand-ctx ns))
-  (define post-expansion-sc (root-expand-context-post-expansion-scope root-ctx))
+  (define post-expansion-sc (post-expansion-scope (root-expand-context-post-expansion root-ctx)))
   (values (seteq post-expansion-sc)
           (set-remove (list->seteq (root-expand-context-module-scopes root-ctx))
                       post-expansion-sc)))

racket/src/expander/expand/body.rkt

@@ -51,8 +51,8 @@
                                  [name #f]
                                  [only-immediate? #t]
                                  [def-ctx-scopes def-ctx-scopes]
-                                 [post-expansion-scope #:parent root-expand-context inside-sc]
+                                 [post-expansion #:parent root-expand-context
-                                 [post-expansion-scope-action add-scope]
+                                                 (lambda (s) (add-scope s inside-sc))]
                                  [scopes (cons inside-sc
                                                (expand-context-scopes ctx))]
                                  [use-site-scopes #:parent root-expand-context (box null)]
@@ -243,7 +243,7 @@
                                             (expand-context-scopes body-ctx))]
                                    [only-immediate? #f]
                                    [def-ctx-scopes #f]
-                                   [post-expansion-scope #:parent root-expand-context #f]))
+                                   [post-expansion #:parent root-expand-context #f]))
   ;; Helper to expand and wrap the ending expressions in `begin`, if needed:
   (define (finish-bodys)
     (define block->list? (null? val-idss))

racket/src/expander/expand/context.rkt

@@ -42,7 +42,6 @@
           phase      ; current expansion phase; must match phase of `namespace`
           namespace  ; namespace for modules and evaluation
           * env        ; environment for local bindings
-          * post-expansion-scope-action ; function to apply with `post-expansion-scope`
           * scopes     ; list of scopes that should be pruned by `quote-syntax`
           * def-ctx-scopes ; #f or box of list of scopes; transformer-created def-ctxes
           * binding-layer ; changed when a binding is nested; to check already-expanded
@@ -76,8 +75,7 @@
   (define root-ctx (namespace-get-root-expand-ctx ns))
   (expand-context (root-expand-context-self-mpi root-ctx)
                   (root-expand-context-module-scopes root-ctx)
-                  (root-expand-context-post-expansion-scope root-ctx)
+                  (root-expand-context-post-expansion root-ctx)
-                  (root-expand-context-post-expansion-shifts root-ctx)
                   (root-expand-context-top-level-bind-scope root-ctx)
                   (root-expand-context-all-scopes-stx root-ctx)
                   (root-expand-context-use-site-scopes root-ctx)
@@ -90,7 +88,6 @@
                   (namespace-phase ns)
                   ns
                   empty-env
-                  push-scope ; post-expansion-scope-action
                   null ; scopes
                   #f   ; def-ctx-scopes [=> don't record scopes to be stipped for `quote-syntax`]
                   (root-expand-context-frame-id root-ctx) ; binding-layer
@@ -121,8 +118,7 @@
   (struct*-copy expand-context ctx
                 [self-mpi #:parent root-expand-context (root-expand-context-self-mpi root-ctx)]
                 [module-scopes #:parent root-expand-context (root-expand-context-module-scopes root-ctx)]
-                [post-expansion-scope #:parent root-expand-context (root-expand-context-post-expansion-scope root-ctx)]
+                [post-expansion #:parent root-expand-context (root-expand-context-post-expansion root-ctx)]
-                [post-expansion-shifts #:parent root-expand-context (root-expand-context-post-expansion-shifts root-ctx)]
                 [top-level-bind-scope #:parent root-expand-context (root-expand-context-top-level-bind-scope root-ctx)]
                 [all-scopes-stx #:parent root-expand-context (root-expand-context-all-scopes-stx root-ctx)]
                 [use-site-scopes #:parent root-expand-context (root-expand-context-use-site-scopes root-ctx)]
@@ -170,7 +166,7 @@
    [else (struct*-copy expand-context ctx
                        [context 'expression]
                        [name #f]
-                       [post-expansion-scope #:parent root-expand-context #f])]))
+                       [post-expansion #:parent root-expand-context #f])]))
 
 ;; Adjusts `ctx` to make it suitable for a non-tail position
 ;; in an `begin` form, possibly in a 'top-level or 'module context

racket/src/expander/expand/definition-context.rkt

@@ -232,9 +232,6 @@
   (define def-ctx-scopes (if (expand-context-def-ctx-scopes ctx)
                              (unbox (expand-context-def-ctx-scopes ctx))
                              null))
-  (define placeholder-sc (and intdefs
-                              (not (null? intdefs))
-                              (new-scope 'macro)))
   (struct*-copy expand-context ctx
                 [context context]
                 [env (add-intdef-bindings (expand-context-env ctx)
@@ -257,29 +254,16 @@
                               ;; Special ID 'all means "use-site scopes for all expansions"
                               'all]
                              [else (or frame-id i-frame-id)]))]
-                [post-expansion-scope
+                [post-expansion #:parent root-expand-context
-                 #:parent root-expand-context
+                                (let ([pe (and same-kind?
-                 (or (and same-kind?
+                                               (or (pair? context)
-                          (memq context '(module module-begin top-level))
+                                                   (memq context '(module module-begin top-level)))
-                          (root-expand-context-post-expansion-scope ctx))
+                                               (root-expand-context-post-expansion ctx))])
-                     ;; Placeholder to make sure `post-expansion-scope-action`
+                                  (cond
-                     ;; is used
+                                    [(and intdefs (not (null? intdefs)))
-                     placeholder-sc)]
+                                     (lambda (s)
-                [post-expansion-shifts
+                                       (add-intdef-scopes (apply-post-expansion pe s) intdefs))]
-                 #:parent root-expand-context
+                                    [else pe]))]
-                 (if (and same-kind?
-                          (eq? context 'top-level))
-                     (root-expand-context-post-expansion-shifts ctx)
-                     null)]
-                [post-expansion-scope-action
-                 (let ([act (expand-context-post-expansion-scope-action ctx)])
-                   (if (and intdefs (not (null? intdefs)))
-                       (lambda (s sc)
-                         (define s2 (if (eq? sc placeholder-sc)
-                                        s
-                                        (act s sc)))
-                         (add-intdef-scopes s2 intdefs))
-                       act))]
                 [scopes
                  (append def-ctx-scopes
                          (expand-context-scopes ctx))]

racket/src/expander/expand/lift-context.rkt

@@ -85,9 +85,10 @@
     ;; Add the namespace's post-expansion scope (i.e., the inside-edge
     ;; scope) so that the binding has a specific phase:
     (define post-scope
-      (root-expand-context-post-expansion-scope
+      (post-expansion-scope
-       (namespace-get-root-expand-ctx
+       (root-expand-context-post-expansion
-        (expand-context-namespace ctx))))
+        (namespace-get-root-expand-ctx
+         (expand-context-namespace ctx)))))
     (define tl-ids (for/list ([id (in-list ids)])
                      (add-scope id post-scope)))
     ;; Bind the identifier:

racket/src/expander/expand/main.rkt

@@ -385,7 +385,7 @@
    (define result-s (flip-scope transformed-s intro-scope))
    ;; In a definition context, we need to add the inside-edge scope to
    ;; any expansion result
-   (define post-s (maybe-add-post-expansion-scope result-s ctx))
+   (define post-s (maybe-add-post-expansion result-s ctx))
    ;; Track expansion:
    (define tracked-s (syntax-track-origin post-s cleaned-s (or origin-id (if (syntax-identifier? s) s (car (syntax-e s))))))
    (define rearmed-s (taint-dispatch tracked-s (lambda (t-s) (syntax-rearm t-s s)) (expand-context-phase ctx)))
@@ -455,20 +455,14 @@
        (or (eq? current-frame-id bind-frame-id)
            (eq? current-frame-id 'all))))
 
-(define (maybe-add-post-expansion-scope s ctx)
+(define (maybe-add-post-expansion s ctx)
-  (cond
+  ;; We may be in a definition context where, say, an inside-edge scope
-   [(root-expand-context-post-expansion-scope ctx)
+  ;; needs to be added to any immediate macro expansion; that way,
-    ;; We're in a definition context where, say, an inside-edge scope
+  ;; if the macro expands to a definition form, the binding will be
-    ;; needs to be added to any immediate macro expansion; that way,
+  ;; in the definition context's scope. The sepcific action depends
-    ;; if the macro expands to a definition form, the binding will be
+  ;; on the expansion context.
-    ;; in the definition context's scope. The sepcific action depends
+  (apply-post-expansion (root-expand-context-post-expansion ctx)
-    ;; on the expansion context.
+                        s))
-    (define new-s
-      ((expand-context-post-expansion-scope-action ctx)
-       s
-       (root-expand-context-post-expansion-scope ctx)))
-    (syntax-add-shifts new-s (root-expand-context-post-expansion-shifts ctx))]
-   [else s]))
 
 (define (accumulate-def-ctx-scopes ctx def-ctx-scopes)
   ;; Move any accumulated definition-context scopes to the `scopes`
@@ -627,7 +621,7 @@
                            (expand-context-stops ctx)
                            empty-free-id-set)]
                 [def-ctx-scopes #f]
-                [post-expansion-scope #:parent root-expand-context #f]))
+                [post-expansion #:parent root-expand-context #f]))
 
 ;; Expand and evaluate `s` as a compile-time expression, ensuring that
 ;; the number of returned values matches the number of target

racket/src/expander/expand/module.rkt

@@ -162,7 +162,7 @@
    (define ctx (struct*-copy expand-context (copy-root-expand-context init-ctx root-ctx)
                              [allow-unbound? #f]
                              [namespace m-ns]
-                             [post-expansion-scope-action add-scope]
+                             [post-expansion #:parent root-expand-context (lambda (s) (add-scope s inside-scope))]
                              [phase phase]
                              [just-once? #f]))
    
@@ -356,7 +356,7 @@
          (define body-ctx (struct*-copy expand-context (accumulate-def-ctx-scopes partial-body-ctx def-ctx-scopes)
                                         [stops empty-free-id-set]
                                         [def-ctx-scopes #f]
-                                        [post-expansion-scope #:parent root-expand-context #f]
+                                        [post-expansion #:parent root-expand-context #f]
                                         [to-module-lifts (make-to-module-lift-context phase
                                                                                       #:shared-module-ends module-ends
                                                                                       #:end-as-expressions? #t)]))
@@ -407,7 +407,7 @@
      
      (define submod-ctx (struct*-copy expand-context ctx
                                       [frame-id #:parent root-expand-context #f]
-                                      [post-expansion-scope #:parent root-expand-context #f]
+                                      [post-expansion #:parent root-expand-context #f]
                                       [namespace submod-m-ns]))
      
      (define declare-enclosing-module
@@ -905,10 +905,10 @@
             (semi-parsed-define-values s syms scoped-ids rhs))))
 
 (define (add-post-expansion-scope bodys ctx)
-  (define sc (root-expand-context-post-expansion-scope ctx))
+  (define pe (root-expand-context-post-expansion ctx))
-  (if sc
+  (if pe
       (for/list ([body (in-list bodys)])
-        (add-scope body sc))
+        (apply-post-expansion pe body))
       bodys))
 
 ;; ----------------------------------------
@@ -1166,8 +1166,8 @@
   (let* ([s (syntax-property s 'module-body-context (root-expand-context-all-scopes-stx root-ctx))]
          [s (syntax-property s
                              'module-body-inside-context
-                             (add-scope empty-syntax
+                             (apply-post-expansion (root-expand-context-post-expansion root-ctx)
-                                        (root-expand-context-post-expansion-scope root-ctx)))])
+                                                   empty-syntax))])
     s))
 
 ;; ----------------------------------------
@@ -1336,7 +1336,7 @@
                    (struct*-copy expand-context ctx
                                  [context 'module]
                                  [stops empty-free-id-set]
-                                 [post-expansion-scope #:parent root-expand-context #f])
+                                 [post-expansion #:parent root-expand-context #f])
                    self
                    #:always-produce-compiled? #t
                    #:keep-enclosing-scope-at-phase keep-enclosing-scope-at-phase

racket/src/expander/expand/root-expand-context.rkt

@@ -9,7 +9,10 @@
 
 (provide (struct*-out root-expand-context)
          make-root-expand-context
-         
+
+         apply-post-expansion
+         post-expansion-scope
+
          root-expand-context-encode-for-module
          root-expand-context-decode-for-module)
 
@@ -20,8 +23,9 @@
 (struct* root-expand-context
          (self-mpi        ; MPI for the enclosing module during compilation
           module-scopes   ; list of scopes for enclosing module or top level; includes next two fields
-          * post-expansion-scope  ; #f or scope to add to every expansion; often module's inside edge
+          * post-expansion  ; #f, a shifted multiscope to push to every expansion (often module's inside edge),
-          * post-expansion-shifts ; a list of MPIshifts to go with `post-expansion-scope`
+          ;                   a pair of a sms and a list of shifts, or a procedure (when not at the actual
+          ;                   root, because an actual root needs to be marshalable)
           top-level-bind-scope  ; #f or a scope to constrain expansion bindings; see "expand-bind-top.rkt"
           all-scopes-stx  ; scopes like the initial import, which correspond to original forms
           * use-site-scopes ; #f or boxed list: scopes that should be pruned from binders
@@ -41,8 +45,7 @@
                                initial-scopes))
   (root-expand-context self-mpi
                        module-scopes
-                       post-expansion-scope
+                       post-expansion-scope ; post-expansion
-                       null ; post-expansion-shifts
                        (new-scope 'module) ; top-level-bind-scope
                        (or all-scopes-stx
                            (add-scopes empty-syntax module-scopes))
@@ -54,13 +57,27 @@
 
 ;; ----------------------------------------
 
+(define (apply-post-expansion pe s)
+  (cond
+    [(not pe) s]
+    [(shifted-multi-scope? pe) (push-scope s pe)]
+    [(pair? pe) (syntax-add-shifts (add-scope s (car pe)) (cdr pe))]
+    [else (pe s)]))
+
+(define (post-expansion-scope pe)
+  (cond
+    [(shifted-multi-scope? pe) pe]
+    [(pair? pe) (car pe)]
+    [else (error 'post-expansion-scope "internal error: cannot extract scope from ~s" pe)]))
+
+;; ----------------------------------------
+
 ;; Encode information in a syntax object that can be serialized and deserialized
 (define (root-expand-context-encode-for-module ctx orig-self new-self)
   (datum->syntax
    #f
    (vector (add-scopes empty-syntax (root-expand-context-module-scopes ctx))
-           (syntax-add-shifts (add-scope empty-syntax (root-expand-context-post-expansion-scope ctx))
+           (apply-post-expansion (root-expand-context-post-expansion ctx) empty-syntax)
-                              (root-expand-context-post-expansion-shifts ctx))
            (syntax-module-path-index-shift (root-expand-context-all-scopes-stx ctx) orig-self new-self)
            (add-scopes empty-syntax (unbox (root-expand-context-use-site-scopes ctx)))
            (for/hasheqv ([(phase ht) (in-hash (root-expand-context-defined-syms ctx))]) ; make immutable
@@ -85,8 +102,8 @@
            vec-s))
   (root-expand-context self
                        (extract-scope-list (vector-ref vec 0)) ; module-scopes
-                       (extract-scope (vector-ref vec 1))      ; post-expansion-scope
+                       (cons (extract-scope (vector-ref vec 1))
-                       (extract-shifts (vector-ref vec 1))     ; post-expansion-scope-shifts
+                             (extract-shifts (vector-ref vec 1))) ; post-expansion
                        (new-scope 'module)                     ; top-level-bind-scope
                        (vector-ref vec 2)                      ; all-scopes-stx
                        (box (extract-scope-list (vector-ref vec 3))) ; use-site-scopes

racket/src/expander/namespace/api.rkt

@@ -58,7 +58,7 @@
   (check who syntax? s)
   (check who namespace? ns)
   (define root-ctx (namespace-get-root-expand-ctx ns))
-  (define post-scope (root-expand-context-post-expansion-scope root-ctx))
+  (define post-scope (post-expansion-scope (root-expand-context-post-expansion root-ctx)))
   (define other-namespace-scopes (for/list ([sc (in-set
                                                  ;; `all-scopes-stx` corresponds to the initial import
                                                  (syntax-scope-set (root-expand-context-all-scopes-stx root-ctx)