(module shared mzscheme
  
  (require "my-macros.ss"
	   "highlight-placeholder.ss"
	   (lib "contracts.ss")
           (lib "etc.ss")
           (lib "list.ss")
           (lib "class.ss")
           (lib "mred.ss" "mred"))

  ; CONTRACTS
  
  (define varref-set? (listof identifier?))
  (define binding-set? (union varref-set? (symbols 'all)))
  (define (arglist? v)
    (or (null? v)
        (identifier? v)
        (and (pair? v)
             ((flat-named-contract-predicate (cons/p identifier? arglist?)) v))        
        (and (syntax? v) 
             ((flat-named-contract-predicate (cons/p identifier? arglist?)) (syntax-e v)))))
  
  (provide/contract
   [varref-set-remove-bindings (-> varref-set? varref-set? varref-set?)]
   [binding-set-varref-set-intersect (-> binding-set? varref-set? binding-set?)]
   [binding-set-union (-> (listof binding-set?) binding-set?)]
   [varref-set-union (-> (listof varref-set?) varref-set?)]
   [in-closure-table (-> any? boolean?)]
   [sublist (-> number? number? list? list?)]
   [attach-info (-> syntax? syntax? syntax?)]
   [transfer-info (-> syntax? syntax? syntax?)]
   [arglist->ilist (-> arglist? any)]
   [arglist-flatten (-> arglist? (listof identifier?))])
  
  (provide
   step-result?
   exp-without-holes?
   exp-with-holes?
   (struct before-after-result (finished-exprs exp redex post-exp reduct after-exprs))
   (struct before-error-result (finished-exprs exp redex err-msg after-exprs))
   (struct error-result (finished-exprs err-msg))
   (struct finished-result (finished-exprs))
   (struct finished-stepping ())
   list-take
   list-partition
   (struct closure-record (name mark constructor? lifted-name))
   *unevaluated* 
   no-sexp
   skipped-step
   struct-flag
   multiple-highlight
   flatten-take
   closure-table-put!
   closure-table-lookup
   insert-highlighted-value
   get-lifted-var
   get-arg-var
   zip
   let-counter
   skipto-annotate
   skipto-reconstruct
   syntax-pair-map
   make-queue ; -> queue
   queue-push ; queue val -> 
   queue-pop ; queue -> val
   queue-length ; queue -> num
   rebuild-stx ; datum syntax-object -> syntax-object
   break-kind? ; predicate
   break-contract ; contract
   varref-set? ; predicate
   binding-set? ; predicate
   ; get-binding-name
   ; bogus-binding?
   if-temp
   ; get-lifted-gensym
   ; expr-read
   ; set-expr-read!
   )
  
  ; an exp-with-holes is either:
  ; - a pair of exp-with-holes's,
  ; - null,
  ; - a symbol, or
  ; - the highlight-placeholder
  
  (define exp-without-holes-base-case? (union symbol? number? string? null? (lambda (v) (is-a? v snip%))))
  
  (define exp-without-holes?
    (union  exp-without-holes-base-case?
           (and/f pair? (cons/p (lx ((flat-named-contract-predicate exp-without-holes?) _))
                                (lx ((flat-named-contract-predicate exp-without-holes?) _))))))
  
  (define exp-with-holes-base-case? 
    (union exp-without-holes-base-case?
          (lx (eq? _ highlight-placeholder))))
  
  (define exp-with-holes?
    (union exp-with-holes-base-case?
          (and/f pair? (cons/p (lx ((flat-named-contract-predicate exp-with-holes?) _)) 
                               (lx ((flat-named-contract-predicate exp-with-holes?) _))))))
  
  ; A step-result is either:
  ; (make-before-after-result finished-exprs exp redex reduct)
  ; or (make-before-error-result finished-exprs exp redex err-msg)
  ; or (make-error-result finished-exprs err-msg)
  ; or (make-finished-result finished-exprs)
  
  (define-struct before-after-result (finished-exprs exp redex post-exp reduct after-exprs))
  (define-struct before-error-result (finished-exprs exp redex err-msg after-exprs))
  (define-struct error-result (finished-exprs err-msg))
  (define-struct finished-result (finished-exprs))
  (define-struct finished-stepping ())
  
  (define step-result? (union before-after-result? before-error-result? error-result? finished-result? finished-stepping?))
  
  ; the closure record is placed in the closure table

  (define-struct closure-record (name mark constructor? lifted-name))

  ; bogus-binding is used so that we can create legal bindings for temporary variables
  
  (define (create-bogus-binding name)
    (let* ([gensymed-name (gensym name)]
           [binding (datum->syntax-object #'here gensymed-name)])
      binding))
  
  ; make-binding-source creates a pool of bindings, indexed by arbitrary keys. These bindings
  ; not eq? to any other bindings[*], but a client can always get the same binding by
  ; invoking the resulting procedure with the same key (numbers work well). make-binding-source
  ; also takes a string which will be part of the printed representation of the binding's
  ; name; this makes debugging easier.
  ; [*] actually, this is not true if you don't use a one-to-one function as the binding-maker
  ; make-gensym-source : (string -> (key -> binding))
  
  (define (make-binding-source id-string binding-maker key-displayer)
    (let ([assoc-table (make-hash-table 'weak)])
      (lambda (key)
        (let ([maybe-fetch (hash-table-get assoc-table key (lambda () #f))])
          (or maybe-fetch
              (begin
                (let* ([new-binding (binding-maker 
                                     (string-append id-string (key-displayer key) "-"))])
                  (hash-table-put! assoc-table key new-binding)
                  new-binding)))))))
  
  
  ; get-arg-var maintains a list of bindings associated with the non-negative
  ; integers.  These symbols are used in the elaboration of applications; the nth
  ; in the application is evaluated and stored in a variable whose name is the nth
  ; gensym supplied by get-arg-var.
  
  (define get-arg-var
    (make-binding-source "arg" create-bogus-binding number->string))
  
  ; test cases: (returns #t on success)
;  (printf "test of get-arg-binding: ~a\n"
;          (let* ([arg3 (get-arg-var 3)]
;                 [arg2 (get-arg-var 2)]
;                 [arg1 (get-arg-var 1)]
;                 [arg2p (get-arg-var 2)])
;            (and (not (eq? arg3 arg2))
;                 (not (eq? arg3 arg1))
;                 (not (eq? arg3 arg2p))
;                 (not (eq? arg2 arg1))
;                 (eq? arg2 arg2p)
;                 (not (eq? arg1 arg2p)))))

  
  ; get-lifted-var maintains the mapping between let-bindings and the syntax object
  ; which is used to capture its index at runtime.
  ; unfortunately, it can't use "make-binding-source" because you need to compare the items 
  ; with module-variable=?, which means that hash tables won't work.
  
  ; my weak-assoc lists are lists of two-element lists, where the first one is in a weak box.
  ; furthermore, the whole thing is in a box, to allow it to be banged when needed.
  
  (define (weak-assoc-add boxed-lst key value)
       (set-box! boxed-lst (cons (list (make-weak-box key) value) (unbox boxed-lst))))
  
  (define (weak-assoc-search boxed-lst key eq-fun)
    (let* ([lst (unbox boxed-lst)]
           [found-val #f]
           [stripped (let loop ([remaining lst])
                       (if (null? remaining)
                           null
                           (let* ([first (car remaining)]
                                  [first-key (weak-box-value (car first))])
                             (if first-key
                                 (if (eq-fun key first-key)
                                     (begin 
                                       (set! found-val (cadr first))
                                       remaining)
                                     (cons first
                                           (loop (cdr remaining))))
                                 (loop (cdr remaining))))))])
      (set-box! boxed-lst stripped)
      found-val))
  
  ; test cases:
  ;  (define wa (box null))
  ;  (define-struct test ())
  ;  (weak-assoc-add wa 3 4)
  ;  (weak-assoc-add wa 9 10)
  ;  (= (weak-assoc-search wa 3 =) 4)
  ;  (= (weak-assoc-search wa 9 =) 10)
  ;  (= (weak-assoc-search wa 3 =) 4)
  ;  (= (length (unbox wa)) 2)
  ;  (define my-struct (make-test))
  ;  (weak-assoc-add wa my-struct 14)
  ;  (= (length (unbox wa)) 3)
  ;  (= (weak-assoc-search wa my-struct eq?) 14)
  ;  (set! my-struct #f)
  ;  (collect-garbage)
  ;  (= (length (unbox wa)) 3)
  ;  (= (weak-assoc-search wa 3 =) 4)
  ;  (= (length (unbox wa)) 2)
  
  (define lifted-index 0)
  (define (next-lifted-symbol str)
    (let ([index lifted-index]) 
      (set! lifted-index (+ lifted-index 1))
      (datum->syntax-object #'here (string->symbol (string-append str (number->string index))))))
 
  (define get-lifted-var
   (let ([assoc-table (box null)])
      (lambda (stx)
        (let ([maybe-fetch (weak-assoc-search assoc-table stx module-identifier=?)])
          (or maybe-fetch
              (begin
                (let* ([new-binding (next-lifted-symbol
                                     (string-append "lifter-" (format "~a" (syntax-object->datum stx)) "-"))])
                  (weak-assoc-add assoc-table stx new-binding)
                  new-binding)))))))
  
  (define if-temp (syntax-property (datum->syntax-object #'here 'if-temp) 'stepper-binding-type 'stepper-temp))

  ; gensyms needed by many modules:

  ; no-sexp is used to indicate no sexpression for display.
  ; e.g., on an error message, there's no sexp.
  (define no-sexp (gensym "no-sexp-"))
  
  ; skipped-step is used to indicate that the "before" step was skipped.
  (define skipped-step (gensym "skipped-step-"))

  ; multiple-highlight is used to indicate multiple highlighted expressions
  (define multiple-highlight (gensym "multiple-highlight-"))
  
  ; *unevaluated* is the value assigned to temps before they are evaluated. It's not a symbol so
  ; it won't need quoting in the source.  Bit of a hack, I know.
  (define-struct *unevaluated-struct* ())
  (define *unevaluated* (make-*unevaluated-struct*))
 
  ; struct-flag : uninterned symbol
  (define struct-flag (gensym "struct-flag-"))
  
  ; unit-result-values-list
  (define unit-result-values-list (gensym "unit-result-vaues-list"))
  
  ; list-partition takes a list and a number, and returns a 2vals containing 2 lists; the first one contains the
  ; first n elements of the list, and the second contains the remainder.  If n is greater than
  ; the length of the list, the exn:application:mismatch exception is raised.
  
  (define (list-partition lst n)
    (if (= n 0)
        (2vals null lst)
        (if (null? lst)
            (list-ref lst 0) ; cheap way to generate exception
            (let*-2vals ([(first rest) (list-partition (cdr lst) (- n 1))])
              (2vals (cons (car lst) first) rest)))))

  ; to perform source correlation, we use the 'register-client' ability of zodiac to
  ; add fields to parsed structures at runtime.
  
;  (define expr-read read-getter)
;  (define set-expr-read! read-setter)
  
  (define (list-take n a-list)
    (if (= n 0)
        null
        (cons (car a-list) (list-take (- n 1) (cdr a-list)))))
  
  (define (flatten-take n a-list)
    (apply append (list-take n a-list)))
  
  (define-values (closure-table-put! closure-table-lookup in-closure-table)
    (let ([closure-table (make-hash-table 'weak)])
      (values
       (lambda (key value)
	 (hash-table-put! closure-table key value)
	 key)                                  ; this return allows a run-time-optimization
       (lambda args ; key or key & failure-thunk
         (apply hash-table-get closure-table args))
       (lambda (key)
         (with-handlers ([exn:application:mismatch? (lambda (exn) #f)])
           (hash-table-get closure-table key)
           #t)))))
  
  ;(begin (closure-table-put! 'foo 'bar)
  ;       (and (eq? (in-closure-table 'blatz) #f)
  ;            (eq? (in-closure-table 'foo) #t)))
 
  ; insert-highlighted-value : sexp sexp -> sexp
  ; replaces highlight-placeholder in the first sexp with the second sexp
  
  (define (insert-highlighted-value exp inserted)
    (let ([recur (lambda (exp) (insert-highlighted-value exp inserted))])
      (cond [(list? exp)
             (map recur exp)]
            [(vector? exp)
             (list->vector (map recur (vector->list exp)))]
            [(eq? exp highlight-placeholder)
             inserted]
            [else exp])))
  
  ;; arglist : for our puposes, an ilist is defined like this:
  ;; arglist : (union identifier? null? (cons identifier? arglist?) (syntax (cons identifier? arglist?))
  ;; ... where an ilist val can be anything _except_ a pair or null

  ;; arglist->ilist : turns an (possibly improper) arglist into a (possibly improper) list of syntax objects

  (define (arglist->ilist arglist)
    (let loop ([ilist arglist])
      (cond [(identifier? ilist)
             ilist]
            [(pair? ilist)
             (cons (car ilist)
                   (loop (cdr ilist)))]
            [(and (syntax? ilist) (pair? (syntax-e ilist)))
             (loop (syntax-e ilist))]
            [(null? ilist) null])))
  
  ; arglist-flatten : produces a list containing the elements of the ilist

  (define (arglist-flatten arglist)
    (let loop ([ilist arglist])
      (cond [(identifier? ilist)
             (cons ilist null)]
            [(null? ilist)
             null]
            [(pair? ilist)
             (cons (car ilist) (loop (cdr ilist)))]
            [(and (syntax? ilist)
                  (pair? (syntax-e ilist)))
             (loop (syntax-e ilist))])))
  
  ; zip : (listof 'a) (listof 'b) (listof 'c) ... -> (listof (list 'a 'b 'c ...))
  ; zip reshuffles lists of items into a list of item-lists.  Look at the contract, okay?
  
  (define zip
    (lambda args
      (apply map list args)))
  
  (define let-counter (syntax-property #'let-counter 'stepper-binding-type 'stepper-temp))
  
 
  ; syntax-pair-map (using the def'ns of the MzScheme docs):
  
  (define (syntax-pair-map pair fn)
    (cons (fn (car pair))
          (cond [(syntax? (cdr pair))
                 (fn (cdr pair))]
                [(pair? (cdr pair))
                 (syntax-pair-map (cdr pair) fn)]
                [(null? (cdr pair))
                 null])))
  
  ; functional update package
  
  (define second-arg (lambda (x y) y))
  
  (define up-mappings
    `((rebuild ((,car ,(lambda (stx new) (cons new (cdr stx))))
                (,cdr ,(lambda (stx new) (cons (car stx) new)))
                (,syntax-e ,(lambda (stx new) (datum->syntax-object stx new stx stx)))))
      (discard ((,car ,second-arg)
                (,cdr ,second-arg)
                (,syntax-e ,second-arg)))))
  
  (define (update fn-list val fn traversal)
    (if (null? fn-list)
        (fn val)
        (let* ([down (car fn-list)]
               [up (cadr (assq down (cadr (assq traversal up-mappings))))])
          (up val (update (cdr fn-list) (down val) fn traversal)))))
 
  
  
  (define skipto-annotate
    (contract 
     (-> (listof procedure?) syntax? (-> syntax? syntax?) syntax?)
     (lambda (fn-list stx annotater)
       (update fn-list stx annotater 'rebuild))
     'skipto-annotate
     'caller))

    ; test cases
;  (equal? (syntax-object->datum (skipto-annotate (list syntax-e car syntax-e cdr cdr car) #'((a b c) (d e f) (g h i)) (lambda (dc) #'foo)))
;          '((a b foo) (d e f) (g h i)))

  
  (define skipto-reconstruct
    (contract
     (-> (listof procedure?) syntax? (-> syntax? any?) any?)
     (lambda (fn-list stx reconstructer)
       (update fn-list stx reconstructer 'discard))
     'skipto-reconstruct
     'caller))
  
  
  (define (make-queue)
    (box null))
  
  (define (queue-push queue new)
    (set-box! queue (append (unbox queue) (list new))))

  (define (queue-pop queue)
    (if (null? (unbox queue))
        (error 'queue-pop "no elements in queue")
        (let ([first (car (unbox queue))])
          (set-box! queue (cdr (unbox queue)))
          first)))
  
  (define (queue-length queue)
    (length (unbox queue)))
  
  (define (rebuild-stx new old)
    (datum->syntax-object old new old old))
  
  (define break-kind?
    (symbols 'normal-break 'result-exp-break 'result-value-break 'double-break 'late-let-break))
    
  (define break-contract
    (-> continuation-mark-set? break-kind? list? any?))
  
  ; BINDING-/VARREF-SET FUNCTIONS
  
  ; note: a BINDING-SET which is not 'all may be used as a VARREF-SET.
  ; this is because they both consist of syntax objects, and a binding
  ; answers true to bound-identifier=? with itself, just like a varref
  ; in the scope of that binding would.
  
  ; binding-set-union: (listof BINDING-SET) -> BINDING-SET
  ; varref-set-union: (listof VARREF-SET) -> VARREF-SET
  
  (define-values (binding-set-union
                  varref-set-union)
    (local ((define (set-pair-union a-set b-set comparator)
              (append (remove* a-set b-set comparator) a-set))
            
            (define (varref-set-pair-union a-set b-set)
              (set-pair-union a-set b-set free-identifier=?))
            
            (define (binding-set-pair-union a-set b-set)
              (cond [(eq? a-set 'all) 'all]
                    [(eq? b-set 'all) 'all]
                    [else (set-pair-union a-set b-set eq?)]))
            
            (define (pair-union->many-union fn)
              (lambda (args)
                (foldl fn null args)))
            
            (define binding-set-union
              (pair-union->many-union binding-set-pair-union))
            
            (define varref-set-union
              (pair-union->many-union varref-set-pair-union)))
      (values binding-set-union
              varref-set-union)))

  ; binding-set-varref-set-intersect : BINDING-SET VARREF-SET -> BINDING-SET
  ; return the subset of varrefs that appear in the bindings
  
  (define (binding-set-varref-set-intersect bindings varrefs)
    (cond [(eq? bindings 'all) varrefs]
          [else (filter (lambda (varref)
                          (ormap (lambda (binding)
                                   (bound-identifier=? binding varref))
                                 bindings))
                        varrefs)]))
  
  ; varref-set-remove-bindings : VARREF-SET (BINDING-SET - 'all) -> VARREF-SET
  ; remove bindings from varrefs
  
  (define (varref-set-remove-bindings varrefs bindings)
    (cond [(eq? bindings 'all)
           (error 'varref-set-remove-bindings "binding-set 'all passed as second argument, first argument was: ~s" varrefs)]
          [else (remove* bindings varrefs bound-identifier=?)]))
  
  ; sublist returns the list beginning with element <begin> and ending just before element <end>.
  ; (-> number? number? list? list?)
  (define (sublist begin end lst) 
    (if (= end 0) 
        null
        (if (= begin 0)
            (cons (car lst)
                  (sublist 0 (- end 1) (cdr lst)))
            (sublist (- begin 1) (- end 1) (cdr lst)))))
  
  
  ; attach-info : SYNTAX-OBJECT SYNTAX-OBJECT -> SYNTAX-OBJECT
  ; attach-info attaches to a generated piece of syntax the origin & source information of another.
  ; we do this so that macro unwinding can tell what reconstructed syntax came from what original syntax
  (define (attach-info stx expr)
    (let* ([it (syntax-property stx 'user-origin (syntax-property expr 'origin))]
           [it (syntax-property it 'user-stepper-hint (syntax-property expr 'stepper-hint))]
           [it (syntax-property it 'user-stepper-else (syntax-property expr 'stepper-else))]
           [it (syntax-property it 'user-source (syntax-source expr))]
           [it (syntax-property it 'user-position (syntax-position expr))])
      it))
  
  (define (transfer-info stx expr)
    (let* ([it (syntax-property stx 'user-origin (syntax-property expr 'user-origin))]
           [it (syntax-property it 'user-stepper-hint (syntax-property expr 'user-stepper-hint))]
           [it (syntax-property it 'user-stepper-else (syntax-property expr 'user-stepper-else))]
           [it (syntax-property it 'user-source (syntax-property expr 'user-source))]
           [it (syntax-property it 'user-position (syntax-property expr 'user-position))])
      it))
  )
  
; test cases
;(require shared)
;(load "/Users/clements/plt/tests/mzscheme/testing.ss")
;
;(define (a sym) 
;  (syntax-object->datum (get-lifted-var sym)))
;(define cd-stx 
;  (datum->syntax-object #f 'cd))
;(test 'lifter-ab-0  a (datum->syntax-object #f 'ab))
;(test 'lifter-cd-1 a cd-stx)
;(test 'lifter-ef-2 a (datum->syntax-object #f 'ef))
;(test 'lifter-cd-1 a cd-stx)
;
;(test '(a b c) map syntax-e (arglist->ilist #'(a b c)))
;(test '(a b c) map syntax-e (arglist->ilist #'(a . (b c))))
;(test 'a syntax-e (arglist->ilist #'a))
;(let ([result (arglist->ilist #' (a b . c))])
;  (test 'a syntax-e (car result))
;  (test 'b syntax-e (cadr result))
;  (test 'c syntax-e (cddr result)))
;(test '(a b c) map syntax-e (arglist-flatten #'(a b c)))
;(test '(a b c) map syntax-e (arglist-flatten #'(a . (b c))))
;(test '(a b c) map syntax-e (arglist-flatten #'(a b . c)))
;(test '(a) map syntax-e (arglist-flatten #'a))
;
;(define (add1 x) (+ x 1))
;(test '(3 4 5) ilist-map add1 '(2 3 4))
;(test '(3 4 . 5) ilist-map add1 '(2 3 . 4))
;
;(test '(2 3 4) ilist-flatten '(2 3 4))
;(test '(2 3 4) ilist-flatten '(2 3 . 4))
;
;(define new-queue (make-queue))
;(test (void) queue-push new-queue 1)
;(test (void) queue-push new-queue 2)
;(test 1 queue-pop new-queue)
;(test (void) queue-push new-queue 3)
;(test 2 queue-pop new-queue)
;(test 3 queue-pop new-queue)
;(err/rt-test (queue-pop new-queue) exn:user?)
