#lang racket/base

(require racket/unit
         syntax/kerncase
         syntax/stx
         syntax/source-syntax
         "private/utils.rkt"
         "errortrace-key.rkt"
         (for-template racket/base "errortrace-key.rkt")
         (for-syntax racket/base)) ; for matching

(define original-stx (make-parameter #f))
(define expanded-stx (make-parameter #f))
(define maybe-undefined (make-parameter #hasheq()))

(provide stacktrace@ stacktrace^ stacktrace-imports^
         stacktrace/annotator-imports^ stacktrace/annotator@
         stacktrace-filter^ stacktrace/annotator/filter@ stacktrace/filter@
         stacktrace/errortrace-annotate^ stacktrace/errortrace-annotate@
         stacktrace/filter/errortrace-annotate@
         original-stx expanded-stx)

(define-signature stacktrace-imports^
  (with-mark
   
   test-coverage-enabled
   test-covered
   initialize-test-coverage-point
   
   profile-key
   profiling-enabled
   initialize-profile-point
   register-profile-start
   register-profile-done))

(define-signature stacktrace/annotator-imports^
  (with-mark
   
   test-coverage-point
   
   profile-key
   profiling-enabled
   initialize-profile-point
   register-profile-start
   register-profile-done))

(define-signature stacktrace-filter^
  (should-annotate?))

(define-signature test-coverage-init^
  (initialize-test-coverage
   ))

;; The intentionally-undocumented format of bindings introduced by `make-st-mark` is:
;; (cons syntax? (cons syntax? srcloc-list))

;; The first syntax object is the original annotated source expression as a (shrunken)
;; datum.

;; The second syntax object is some part of the original syntax as a (shrunken)
;; datum, which contains the code that expanded to the annotated expression.

(define-signature stacktrace^
  (annotate-top
   annotate
   make-st-mark
   st-mark-source
   st-mark-bindings))

(define-signature errortrace-annotate^
  (errortrace-annotate))

(define-signature stacktrace/errortrace-annotate^
  ((open stacktrace^)
   (open errortrace-annotate^)))

(define base-phase
  (variable-reference->module-base-phase (#%variable-reference)))

(define orig-inspector (variable-reference->module-declaration-inspector
                        (#%variable-reference)))

(define (rearm orig new)
  (syntax-rearm new orig))

(define (disarm orig)
  (syntax-disarm orig orig-inspector))

(define-unit stacktrace/annotator/filter@
  (import stacktrace/annotator-imports^
          stacktrace-filter^)
  (export stacktrace^)

  (define (short-version v depth)
    (cond
      [(identifier? v) (syntax-e v)]
      [(null? v) null]
      [(vector? v) (if (zero? depth)
                       #(....)
                       (list->vector
                        (short-version (vector->list v) (sub1 depth))))]
      [(box? v) (if (zero? depth)
                    #&(....)
                    (box (short-version (unbox v) (sub1 depth))))]
      [(pair? v)
       (cond
         [(zero? depth) '(....)]
         [(memq (syntax-e (car v)) '(#%app #%top))
          (short-version (cdr v) depth)]
         [else
          (cons (short-version (car v) (sub1 depth))
                (short-version (cdr v) (sub1 depth)))])]
      [(syntax? v) (short-version (syntax-e v) depth)]
      [else v]))
  
  (define current-recover-table (make-parameter #f))

  (define (make-st-mark stx phase)
    (unless (syntax? stx)
      (error 'make-st-mark
             "expected syntax object as argument, got ~e" stx))
    (cond
      [(should-annotate? stx phase)
       ;; this horrible indirection is needed because the errortrace
       ;; unit is invoked only once but annotate-top might be called
       ;; many times with diferent values for original-stx and
       ;; expanded-stx
       (define recover-table (current-recover-table))
       (define (make-new-recover)
         (recover-source-syntax (original-stx) (expanded-stx)
                                #:traverse-now? #t))
       (define recover (if recover-table
                           (hash-ref! recover-table 
                                      (cons (original-stx) (expanded-stx))
                                      make-new-recover)
                           (make-new-recover)))
       (define better-stx (or (and stx (recover stx)) stx))
       (with-syntax ([quote (syntax-shift-phase-level #'quote phase)])
         #`(quote (#,(short-version better-stx 10)
                   #,(syntax-source stx)
                   #,(syntax-line stx)
                   #,(syntax-column stx)
                   #,(syntax-position stx)
                   #,(syntax-span stx))))]
      [else #f]))
  
  (define (st-mark-source src)
    (and src
         (datum->syntax #f (car src) (cdr src) #f)))
  
  (define (st-mark-bindings x) null)
  
  
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  ;; Profiling instrumenter
  
  ;; profile-point :
  ;;   (syntax[list of exprs] symbol-or-#f syntax boolean
  ;;    -> syntax[list of exprs])
  
  ;; This procedure is called by `annotate' and `annotate-top' to wrap
  ;; expressions with profile collecting information.  Returning the
  ;; first argument means no profiling information is collected.
  
  ;; The second argument is the point's inferred name, if any, the third
  ;; argument is the source expression, and the fourth argument is #t for
  ;; a transformer expression and #f for a normal expression.
  
  (define (profile-point bodies name expr phase)
    (let ([key (gensym 'profile-point)])
      (initialize-profile-point key name expr)
      (with-syntax ([key (datum->syntax #f key (quote-syntax here))]
                    [start (datum->syntax
                            #f (gensym) (quote-syntax here))]
                    [profile-key (datum->syntax
                                  #f profile-key (quote-syntax here))]
                    [register-profile-start register-profile-start]
                    [register-profile-done register-profile-done]
                    [app (syntax-shift-phase-level #'#%plain-app (- phase base-phase))]
                    [lt (syntax-shift-phase-level #'let-values (- phase base-phase))]
                    [qt (syntax-shift-phase-level #'quote (- phase base-phase))]
                    [bgn (syntax-shift-phase-level #'begin (- phase base-phase))]
                    [wcm (syntax-shift-phase-level #'with-continuation-mark (- phase base-phase))])
        (with-syntax ([rest
                       (insert-at-tail*
                        (syntax (app (qt register-profile-done) (qt key) start))
                        bodies
                        phase)])
          (syntax
           (lt ([(start) (app (qt register-profile-start) (qt key))])
             (wcm 
              (qt profile-key)
              (qt key)
              (bgn . rest))))))))
  
  (define (insert-at-tail* e exprs phase)
    (let ([new
           (rebuild exprs
                    (let loop ([exprs exprs])
                      (if (stx-null? (stx-cdr exprs))
                          (list (cons (stx-car exprs)
                                      (insert-at-tail
                                       e (stx-car exprs) phase)))
                          (loop (stx-cdr exprs)))))])
      (if (syntax? exprs)
          (rearm exprs new)
          new)))
  
  (define (insert-at-tail se sexpr phase)
    (with-syntax ([expr sexpr]
                  [e se]
                  [bgn (syntax-shift-phase-level #'begin (- phase base-phase))]
                  [bgn0 (syntax-shift-phase-level #'begin0 (- phase base-phase))])
      (kernel-syntax-case/phase sexpr phase
        ;; negligible time to eval
        [id
         (identifier? sexpr)
         (syntax (bgn e expr))]
        [(quote _) (syntax (bgn e expr))]
        [(quote-syntax . _) (syntax (bgn e expr))]
        [(#%top . d) (syntax (bgn e expr))]
        [(#%variable-reference . d) (syntax (bgn e expr))]
        
        ;; No tail effect, and we want to account for the time
        [(#%plain-lambda . _) (syntax (bgn0 expr e))]
        [(case-lambda . _) (syntax (bgn0 expr e))]
        [(set! . _) (syntax (bgn0 expr e))]
        
        [(let-values bindings . body)
         (insert-at-tail* se sexpr phase)]
        [(letrec-values bindings . body)
         (insert-at-tail* se sexpr phase)]
        [(letrec-syntaxes+values sbindings bindings . body)
         (insert-at-tail* se sexpr phase)]
        
        [(begin . _)
         (insert-at-tail* se sexpr phase)]
        [(with-continuation-mark . _)
         (insert-at-tail* se sexpr phase)]
        
        [(begin0 body ...)
         (rearm sexpr (syntax (bgn0 body ... e)))]
        
        [(if test then else)
         ;; WARNING: se inserted twice!
         (rearm
          sexpr
          (rebuild
           sexpr
           (list
            (cons #'then (insert-at-tail se (syntax then) phase))
            (cons #'else (insert-at-tail se (syntax else) phase)))))]
        
        [(#%plain-app . rest)
         (if (stx-null? (syntax rest))
             ;; null constant
             (syntax (bgn e expr))
             ;; application; exploit guaranteed left-to-right evaluation
             (insert-at-tail* se sexpr phase))]

        [(#%expression e)
         (rearm
          sexpr
          (rebuild sexpr
                   (list (cons #'e (insert-at-tail se (syntax e) phase)))))]
        
        [_else
         (error 'errortrace
                "unrecognized (non-top-level) expression form: ~.s"
                (syntax->datum sexpr))])))
  
  (define (profile-annotate-lambda name expr clause bodys-stx phase)
    (let* ([bodys (stx->list bodys-stx)]
           [bodyl (map (lambda (e) (no-cache-annotate e phase))
                       bodys)])
      (rebuild clause
               (if (profiling-enabled)
                   (let ([prof-expr
                          (profile-point bodyl name expr phase)])
                     ;; Tell rebuild to replace first expressions with
                     ;; (void), and replace the last expression with
                     ;; prof-expr:
                     (let loop ([bodys bodys])
                       (if (null? (cdr bodys))
                           (list (cons (car bodys) prof-expr))
                           (cons (cons (car bodys) #'(#%plain-app void))
                                 (loop (cdr bodys))))))
                   ;; Map 1-to-1:
                   (map cons bodys bodyl)))))
  
  (define (keep-lambda-properties orig new)
    (let ([p (syntax-property orig 'method-arity-error)]
          [p2 (syntax-property orig 'inferred-name)])
      (let ([new (if p
                   (syntax-property new 'method-arity-error p)
                   new)])
        (if p2
            (syntax-property new 'inferred-name p2)
            new))))
  
  (define (annotate-let expr phase varss-stx rhss-stx bodys-stx letrec?)
    (let ([varss (syntax->list varss-stx)]
          [rhss  (syntax->list rhss-stx)]
          [bodys (syntax->list bodys-stx)])
      (parameterize ([maybe-undefined (if (and letrec?
                                               (not (andmap (simple-rhs? phase) rhss)))
                                          (add-identifiers
                                           (apply append (map syntax->list varss))
                                           (maybe-undefined))
                                          (maybe-undefined))])
        (let ([rhsl (map
                     (lambda (vars rhs)
                       (define rhs* (no-cache-annotate-named
                                     (syntax-case vars () [(id) (syntax id)] [_else #f])
                                     rhs
                                     phase))
                       (cover-names vars rhs* phase))
                     varss
                     rhss)]
              [bodyl (map (lambda (body) (no-cache-annotate body phase))
                          bodys)])
          (rebuild expr (append (map cons bodys bodyl)
                                (map cons rhss rhsl)))))))


  (define (cover-names names-stx body phase)
    (for/fold ([body body]) ([name (in-list (syntax->list names-stx))])
      (test-coverage-point body name phase)))

  (define ((simple-rhs? phase) expr)
    (kernel-syntax-case/phase expr phase
      [(quote _) #t]
      [(quote-syntax . _) #t]
      [(#%plain-lambda . _) #t]
      [(case-lambda . _) #t]
      [_else #f]))

  (define (add-identifiers ids ht)
    (for/fold ([ht ht]) ([id (in-list ids)])
      (define l (hash-ref ht (syntax-e id) null))
      (hash-set ht (syntax-e id) (cons id l))))

  (define (maybe-undefined? id phase)
    (define l (hash-ref (maybe-undefined) (syntax-e id) null))
    (for/or ([mu-id (in-list l)])
      (free-identifier=? mu-id id phase)))
  
  (define (annotate-seq expr bodys-stx annotate phase)
    (let* ([bodys (syntax->list bodys-stx)]
           [bodyl (map (lambda (b) (annotate b phase)) bodys)])
      (rebuild expr (map cons bodys bodyl))))
  
  (define (rebuild expr replacements)
    (let loop ([expr expr] [same-k (lambda () expr)] [diff-k (lambda (x) x)])
      (let ([a (assq expr replacements)])
        (cond
          [a (diff-k (cdr a))]
          [(pair? expr)
           (loop (car expr)
                 (lambda ()
                   (loop (cdr expr) same-k
                         (lambda (y) (diff-k (cons (car expr) y)))))
                 (lambda (x)
                   (loop (cdr expr)
                         (lambda () (diff-k (cons x (cdr expr))))
                         (lambda (y) (diff-k (cons x y))))))]
          [(vector? expr)
           (loop (vector->list expr) same-k
                 (lambda (x) (diff-k (list->vector x))))]
          [(box? expr)
           (loop (unbox expr) same-k (lambda (x) (diff-k (box x))))]
          [(syntax? expr)
           (if (identifier? expr)
             (same-k)
             (loop (syntax-e expr) same-k
                   (lambda (x) (diff-k (datum->syntax expr x expr expr)))))]
          [else (same-k)]))))
  
  (define (append-rebuild expr end)
    (cond
      [(syntax? expr)
       (datum->syntax expr
                      (append-rebuild (syntax-e expr) end)
                      expr
                      expr)]
      [(pair? expr)
       (cons (car expr) (append-rebuild (cdr expr) end))]
      [(null? expr)
       (list end)]
      [else
       (error 'append-rebuild "shouldn't get here")]))
  
  (define (one-name names-stx)
    (let ([l (syntax->list names-stx)])
      (and (pair? l)
           (null? (cdr l))
           (car l))))
  
  (define (make-annotate top? name)
    (lambda (expr phase)
      (define disarmed-expr (disarm expr))
      (define (with-mrk* mark expr)
        (with-mark mark expr phase))
      (test-coverage-point
       (kernel-syntax-case/phase disarmed-expr phase
         [_
          (identifier? expr)
          (let ([b (identifier-binding expr phase)])
            (cond
             [(eq? 'lexical b)
              ;; lexical variable; a use-before-defined error is possible
              ;; for letrec-bound variables
              (if (maybe-undefined? expr phase)
                  (with-mrk* expr expr)
                  expr)]
             [(and (pair? b) (let-values ([(base rel) (module-path-index-split (car b))])
                               (or base rel)))
              ;; from another module -- no error possible
              expr]
             [else
              ;; might be undefined/uninitialized
              (with-mrk* expr expr)]))]
         
         [(#%top . id)
          ;; might be undefined/uninitialized
          (with-mrk* expr expr)]
         [(#%variable-reference . _)
          ;; no error possible
          expr]
         
         [(define-values names rhs)
          top?
          ;; Can't put annotation on the outside
          (let* ([marked 
                  (with-mrk* expr
                             (no-cache-annotate-named
                              (one-name #'names)
                              (syntax rhs)
                              phase))]
                 [with-coverage (cover-names #'names marked phase)])
            (rearm
             expr
             (rebuild 
              disarmed-expr 
              (list (cons #'rhs with-coverage)))))]
         [(begin . exprs)
          top?
          (rearm
           expr
           (annotate-seq disarmed-expr
                         (syntax exprs)
                         no-cache-annotate-top phase))]
         [(define-syntaxes (name ...) rhs)
          top?
          (let* ([marked (with-mark expr
                           (no-cache-annotate-named
                            (one-name #'(name ...))
                            (syntax rhs)
                            (add1 phase))
                           (add1 phase))]
                 ;; cover at THIS phase, since thats where its bound
                 [rebuilt (rebuild disarmed-expr (list (cons #'rhs marked)))]
                 [with-coverage (cover-names #'(name ...) rebuilt phase)])
            (rearm expr with-coverage))]
         
         [(begin-for-syntax . exprs)
          top?
          (rearm
           expr
           (annotate-seq disarmed-expr
                         (syntax exprs)
                         no-cache-annotate-top 
                         (add1 phase)))]

         [(module name init-import mb)
          (annotate-module expr disarmed-expr 0)]
         [(module* name init-import mb)
          (annotate-module expr disarmed-expr (if (syntax-e #'init-import) 0 phase))]
         
         [(#%expression e)
          (rearm expr #`(#%expression #,(no-cache-annotate (syntax e) phase)))]
         
         ;; No way to wrap
         [(#%require i ...) expr]
         ;; No error possible (and no way to wrap)
         [(#%provide i ...) expr]
         [(#%declare i ...) expr]
         
         ;; No error possible
         [(quote _)
          expr]
         [(quote-syntax . _)
          expr]
         
         ;; Wrap body, also a profile point
         [(#%plain-lambda args . body)
          (rearm
           expr
           (keep-lambda-properties
            expr
            (profile-annotate-lambda name expr disarmed-expr (syntax body)
                                     phase)))]
         [(case-lambda clause ...)
          (with-syntax ([([args . body] ...)
                         (syntax (clause ...))])
            (let* ([clauses (syntax->list (syntax (clause ...)))]
                   [clausel (map
                             (lambda (body clause)
                               (profile-annotate-lambda
                                name expr clause body phase))
                             (syntax->list (syntax (body ...)))
                             clauses)])
              (rearm
               expr
               (keep-lambda-properties
                expr
                (rebuild disarmed-expr (map cons clauses clausel))))))]
         
         ;; Wrap RHSs and body
         [(let-values ([vars rhs] ...) . body)
          (with-mrk* expr
                     (rearm
                      expr
                      (annotate-let disarmed-expr phase
                                    (syntax (vars ...))
                                    (syntax (rhs ...))
                                    (syntax body)
                                    #f)))]
         [(letrec-values ([vars rhs] ...) . body)
          (let ([fm (rearm
                     expr
                     (annotate-let disarmed-expr phase
                                   (syntax (vars ...))
                                   (syntax (rhs ...))
                                   (syntax body)
                                   #t))])
            (kernel-syntax-case/phase expr phase
              [(lv ([(var1) (#%plain-lambda . _)]) var2)
               (and (identifier? #'var2)
                    (free-identifier=? #'var1 #'var2))
               fm]
              [_
               (with-mrk* expr fm)]))]
         ;; This case is needed for `#lang errortrace ...', which uses
         ;; `local-expand' on the module body.
         [(letrec-syntaxes+values sbindings ([vars rhs] ...) . body)
          (let ([fm (rearm
                     expr
                     (annotate-let disarmed-expr phase
                                   (syntax (vars ...))
                                   (syntax (rhs ...))
                                   (syntax body)
                                   #t))])
            (with-mrk* expr fm))]

         ;; Wrap RHS
         [(set! var rhs)
          (let ([new-rhs (no-cache-annotate-named
                          (syntax var)
                          (syntax rhs)
                          phase)])
            ;; set! might fail on undefined variable, or too many values:
            (with-mrk* expr
                       (rearm
                        expr
                        (rebuild disarmed-expr (list (cons #'rhs new-rhs))))))]
         
         ;; Wrap subexpressions only
         [(begin e)
          ;; Single expression: no mark
          (rearm
           expr
           #`(begin #,(no-cache-annotate (syntax e) phase)))]
         [(begin . body)
          (with-mrk* expr
                     (rearm
                      expr
                      (annotate-seq disarmed-expr #'body no-cache-annotate phase)))]
         [(begin0 . body)
          (with-mrk* expr
                     (rearm
                      expr
                      (annotate-seq disarmed-expr #'body no-cache-annotate phase)))]
         [(if tst thn els)
          (let ([w-tst (no-cache-annotate (syntax tst) phase)]
                [w-thn (no-cache-annotate (syntax thn) phase)]
                [w-els (no-cache-annotate (syntax els) phase)])
            (with-mrk* expr
                       (rearm
                        expr
                        (rebuild disarmed-expr (list (cons #'tst w-tst)
                                                     (cons #'thn w-thn)
                                                     (cons #'els w-els))))))]
         [(with-continuation-mark . body)
          (with-mrk* expr
                     (rearm
                      expr
                      (annotate-seq disarmed-expr (syntax body)
                                    no-cache-annotate phase)))]
         
         ;; Wrap whole application, plus subexpressions
         [(#%plain-app . body)
          (cond
           [(stx-null? (syntax body))
            ;; It's a null:
            expr]
           ;; Wrong idea:
           ;;  check for functions that are known to always succeed,
           ;;  in which case we can skip the wrapper:
           ;; Why it's wrong:
           ;;  although these functions always succeed, the argument
           ;;  expressions may not, and we want to include this call
           ;;  in the stack trace if the argument goes wrong
           #;
           [(syntax-case* expr (void cons mcons list list* vector box
                                     vector-immutable)
                          (lambda (a b)
                            (free-identifier=? a b phase base-phase))
              [(_ void . _) #t]
              [(_ cons _ _) #t]
              [(_ mcons _ _) #t]
              [(_ list . _) #t]
              [(_ list* _ . _) #t]
              [(_ vector . _) #t]
              [(_ vector-immutable . _) #t]
              [(_ box _) #t]
              [_else #f])
            (rearm
             expr
             (annotate-seq disarmed-expr (syntax body)
                           no-cache-annotate phase))]
           ;; general case:
           [else
            (with-mrk* expr (rearm
                             expr
                             (annotate-seq disarmed-expr (syntax body)
                                           no-cache-annotate phase)))])]
         
         [_else
          (error 'errortrace "unrecognized expression form~a~a: ~.s"
                 (if top? " at top-level" "")
                 (if (zero? phase) "" (format " at phase ~a" phase))
                 (syntax->datum expr))])
       expr
       phase)))

  (define (annotate-module expr disarmed-expr phase)
    (define shifted-disarmed-expr
      (syntax-shift-phase-level disarmed-expr (- phase)))
    (syntax-case shifted-disarmed-expr ()
      [(mod name init-import mb)
       (syntax-case (disarm #'mb) ()
         [(__plain-module-begin body ...)
          ;; Just wrap body expressions
          (let ([bodys (syntax->list (syntax (body ...)))])
            (let ([bodyl (map (lambda (b)
                                (no-cache-annotate-top b 0))
                              bodys)]
                  [mb #'mb])
              (rearm
               expr
               (syntax-shift-phase-level
                (rebuild
                 shifted-disarmed-expr
                 (list (cons
                        mb
                        (rearm
                         mb
                         (rebuild mb (map cons bodys bodyl))))))
                phase))))])]))
  
  (define no-cache-annotate (make-annotate #f #f))
  (define (annotate expr phase)
    (parameterize ([current-recover-table (make-hash)])
      (no-cache-annotate expr phase)))
  (define no-cache-annotate-top (make-annotate #t #f))
  (define (annotate-top expr phase)
    (parameterize ([current-recover-table (make-hash)])
      (no-cache-annotate-top expr phase)))
  (define (no-cache-annotate-named name expr phase)
    ((make-annotate #f name) expr phase))
  (define (annotate-named name expr phase)
    (parameterize ([current-recover-table (make-hash)])
      (no-cache-annotate-named name expr phase))))

(define-unit stacktrace/annotator/filter/errortrace-annotate@
  (import stacktrace^)
  (export errortrace-annotate^)

  (define (errortrace-annotate top-e [in-compile-handler? #t])

    (define (do-expand e)
      (expand-syntax (add-annotate-property e)))

    (define (do-annotate e expanded-e)
      (parameterize ([original-stx e]
                     [expanded-stx expanded-e])
        (annotate-top expanded-e (namespace-base-phase))))

    (syntax-case top-e ()
      [(mod name . reste)
       (and (identifier? #'mod)
            (free-identifier=? #'mod
                               (namespace-module-identifier)
                               (namespace-base-phase)))
       (if (eq? (syntax-e #'name) 'errortrace-key)
           top-e
           (let ([expanded-e (do-expand top-e)])
             (cond
               [(has-cross-phase-declare?
                 (syntax-case expanded-e ()
                   [(mod name init-import mb) #'mb]))
                expanded-e]
               [else
                (transform-all-modules
                 (do-annotate top-e expanded-e)
                 (lambda (top-e mod-id)
                   (syntax-case top-e ()
                     [(mod name init-import mb)
                      (syntax-case (disarm #'mb) (#%plain-module-begin)
                        [(#%plain-module-begin body ...)
                         (let ([meta-depth ((count-meta-levels 0) #'(begin body ...))])
                           (copy-props
                            top-e
                            #`(#,mod-id name init-import
                                        #,(syntax-rearm
                                           #`(#%plain-module-begin
                                              #,(generate-key-imports meta-depth)
                                              body ...)
                                           #'mb))))])])))])))]
      [_else
       (let ()
         (define e (do-annotate top-e (do-expand top-e)))
         (define meta-depth ((count-meta-levels 0) e))
         (when in-compile-handler?
           ;; We need to force the `require`s now, so that `e` can be compiled.
           ;; It doesn't work to reply on `begin` unrolling for a top-level `eval`,
           ;; because we're in a compile handler and already committed to a single form.
           (for ([i (in-range meta-depth)])
             (namespace-require `(for-meta ,(add1 i) errortrace/errortrace-key))))
         #`(begin
             #,(generate-key-imports-at-phase meta-depth (namespace-base-phase))
             #,e))]))

  (define (has-cross-phase-declare? e)
    (for/or ([a (in-list (or (syntax->list e) '()))])
      (syntax-case* a (#%declare begin) (lambda (a b)
                                          (free-identifier=? a b 0 base-phase))
        [(#%declare kw ...)
         (for/or ([kw (in-list (syntax->list #'(kw ...)))])
           (eq? (syntax-e kw) '#:cross-phase-persistent))]
        [(begin . e)
         (has-cross-phase-declare? #'e)]
        [_ #f])))

  ;; Add the 'errortrace:annotate property everywhere in an original syntax
  ;; object, so that we can recognize pieces from the original program
  ;; after expansion:
  (define (add-annotate-property s)
    (cond
      [(syntax? s)
       (define new-s (syntax-rearm
                      (let ([s (disarm s)])
                        (datum->syntax s
                                       (add-annotate-property (syntax-e s))
                                       s
                                       s))
                      s))
       ;; We could check for "original" syntax here, because partial expansion
       ;; of top-level forms happens before the compile handler gets
       ;; control. As a compromise for programs that may use errortrace
       ;; with `eval` and S-expressions or non-`syntax-original?` arrangements, we add
       ;; the property everywhere:
       (if #t ; (syntax-original? s)
           (syntax-property new-s
                            'errortrace:annotate #t
                            ;; preserve the property in bytecode:
                            #t)
           new-s)]
      [(pair? s)
       (cons (add-annotate-property (car s))
             (add-annotate-property (cdr s)))]
      [(vector? s)
       (for/vector #:length (vector-length s) ([e (in-vector s)])
         (add-annotate-property e))]
      [(box? s) (box (add-annotate-property (unbox s)))]
      [(prefab-struct-key s)
       => (lambda (k)
            (apply make-prefab-struct
                   k
                   (add-annotate-property (cdr (vector->list (struct->vector s))))))]
      [(and (hash? s) (immutable? s))
       (cond
         [(hash-eq? s)
          (for/hasheq ([(k v) (in-hash s)])
            (values k (add-annotate-property v)))]
         [(hash-eqv? s)
          (for/hasheqv ([(k v) (in-hash s)])
            (values k (add-annotate-property v)))]
         [else
          (for/hash ([(k v) (in-hash s)])
            (values k (add-annotate-property v)))])]
      [else s]))

  (define (transform-all-modules stx proc [in-mod-id (namespace-module-identifier)])
    (syntax-case stx ()
      [(mod name init-import mb)
       (syntax-case (disarm #'mb) (#%plain-module-begin)
         [(#%plain-module-begin body ...)
          (let ()
            (define ((handle-top-form phase) expr)
              (define disarmed-expr (disarm expr))
              (syntax-case* disarmed-expr (begin-for-syntax module module*)
                (lambda (a b)
                  (free-identifier=? a b phase 0))
                [(begin-for-syntax body ...)
                 (syntax-rearm
                  #`(#,(car (syntax-e disarmed-expr))
                     #,@(map (handle-top-form (add1 phase))
                             (syntax->list #'(body ...))))
                  expr)]
                [(module . _)
                 (syntax-rearm
                  (transform-all-modules disarmed-expr proc #f)
                  expr)]
                [(module* name init-import . _)
                 (let ([shift (if (syntax-e #'init-import)
                                  0
                                  phase)])
                   (syntax-rearm
                    (syntax-shift-phase-level
                     (transform-all-modules (syntax-shift-phase-level disarmed-expr (- shift)) proc #f)
                     shift)
                    expr))]
                [else expr]))
            (define mod-id (or in-mod-id #'mod))
            (proc
             (copy-props
              stx
              #`(#,mod-id name init-import
                          #,(syntax-rearm
                             #`(#%plain-module-begin
                                . #,(map (handle-top-form 0) (syntax->list #'(body ...))))
                             #'mb)))
             mod-id))])]))

  (define (copy-props orig new)
    (datum->syntax orig (syntax-e new) orig orig)))

;; --------------------------------------------------

(define-unit test-coverage-point@
  (import (prefix in: stacktrace-imports^)
          stacktrace-filter^)
  (export stacktrace/annotator-imports^)
  
  (define with-mark in:with-mark)
  (define profile-key in:profile-key)
  (define profiling-enabled in:profiling-enabled)
  (define initialize-profile-point in:initialize-profile-point)
  (define register-profile-start in:register-profile-start)
  (define register-profile-done in:register-profile-done)
  
  ;; ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  ;; Test case coverage instrumenter
 
  ;; The next procedure is called by `annotate' and `annotate-top' to wrap
  ;; expressions with test suite coverage information.  Returning the
  ;; first argument means no tests coverage information is collected.
 
  ;; test-coverage-point : syntax syntax phase -> syntax
  ;; sets a test coverage point for a single expression
  (define (test-coverage-point body expr phase)
    (if (and (in:test-coverage-enabled)
             (zero? phase) 
             (should-annotate? expr phase))
      (begin (in:initialize-test-coverage-point expr)
             (let ([thunk (in:test-covered expr)])
               (cond [(procedure? thunk)
                      (with-syntax ([body body] [thunk thunk])
                        #'(begin (#%plain-app thunk) body))]
                     [(syntax? thunk)
                      (with-syntax ([body body] [thunk thunk])
                        #'(begin thunk body))]
                     [else body])))
      body)))

;; --------------------------------------------------

(define-unit annotate-if-source@
  (import)
  (export stacktrace-filter^)
  (define (should-annotate? s phase)
    (syntax-source s)))

;; --------------------------------------------------

(define-compound-unit stacktrace/annotator@
  (import (in : stacktrace/annotator-imports^))
  (export out)
  (link [((annotate-if-source : stacktrace-filter^)) annotate-if-source@]
        [((out : stacktrace^)) stacktrace/annotator/filter@ in annotate-if-source]))

;; --------------------------------------------------
  
(define-compound-unit stacktrace/filter@
  (import (in : stacktrace-imports^)
          (filter : stacktrace-filter^))
  (export out)
  (link [((coverage-point : stacktrace/annotator-imports^)) test-coverage-point@ in filter]
        [((out : stacktrace^)) stacktrace/annotator/filter@ coverage-point filter]))

;; --------------------------------------------------
  
(define-compound-unit stacktrace@
  (import (in : stacktrace-imports^))
  (export out)
  (link [((annotate-if-source : stacktrace-filter^)) annotate-if-source@]
        [((coverage-point : stacktrace/annotator-imports^)) test-coverage-point@ in annotate-if-source]
        [((out : stacktrace^)) stacktrace/annotator/filter@ coverage-point annotate-if-source]))

(define-compound-unit stacktrace/filter/errortrace-annotate@
  (import (in : stacktrace-imports^)
          (filter : stacktrace-filter^))
  (export out1 out2)
  (link [((coverage-point : stacktrace/annotator-imports^)) test-coverage-point@ in filter]
        [((out1 : stacktrace^)) stacktrace/annotator/filter@ coverage-point filter]
        [((out2 : errortrace-annotate^)) stacktrace/annotator/filter/errortrace-annotate@ out1]))

(define-compound-unit stacktrace/errortrace-annotate@
  (import (in : stacktrace/annotator-imports^))
  (export out1 out2)
  (link [((annotate-if-source : stacktrace-filter^)) annotate-if-source@]
        [((out1 : stacktrace^)) stacktrace/annotator/filter@ in annotate-if-source]
        [((out2 : errortrace-annotate^)) stacktrace/annotator/filter/errortrace-annotate@ out1]))
