#lang scheme
(require (planet cce/scheme:4:1/planet)
         (for-syntax (planet cce/scheme:4:1/planet)))
(require (prefix-in scheme: scheme)
         (this-package-in private/command-line)
         (for-syntax (this-package-in private/command-line))
         (this-package-in private/collector-exports)
         (this-package-in private/gc-core)
         scheme/gui/dynamic
         (only-in (this-package-in test-harness) generic-test halt-on-errors print-only-errors)
         (for-syntax scheme)
         (for-syntax (this-package-in private/gc-transformer)))

(provide deref
         if and or cond case #%top define define-values let-values let* lambda set!
         else begin
         #%plain-app
         #%require
         test/location=? test/value=?
         (rename-out
          (mutator-halt-on-errors halt-on-errors)
          (mutator-print-only-errors print-only-errors))
         (rename-out
          (mutator-let let)
          (mutator-app #%app)
          (mutator-datum #%datum)
          (collector:cons cons)
          (collector:first first)
          (collector:rest rest)
          (mutator-quote quote)
          (mutator-top-interaction #%top-interaction)
          (mutator-module-begin #%module-begin)))

; Workaround for named-let bug
(define-syntax mutator-let
  (syntax-rules ()
    [(_ ([id expr] ...) body ...)
     (let ([id expr] ...) body ...)]))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Providing lifted procedures


(define-for-syntax required-allocator-stx false)

(define-syntax (mutator-halt-on-errors stx)
  (syntax-case stx ()
    [(_) #'(do-not-annotate (halt-on-errors))]
    [(_ arg) #'(do-not-annotate (#%app halt-on-errors (#%datum . arg)))]))

(define-syntax (mutator-print-only-errors stx)
  (syntax-case stx ()
    [(_) #'(do-not-annotate (print-only-errors))]
    [(_ arg) #'(do-not-annotate (#%app print-only-errors (#%datum . arg)))]))

(provide import-gc)
(define-syntax (import-gc stx)
  (syntax-case stx ()
    [(_ id ...)
     (andmap identifier? (syntax->list #'(id ...)))
     #`(begin
         (when (gc-disable-import-gc?)
           (error 'import-gc "import-gc is disabled"))
         (require (only-in #,required-allocator-stx id ...)))]
    [(_ maybe-id ...) 
     (ormap (λ (v) (and (not (identifier? v)) v)) (syntax->list #'(maybe-id ...)))
     (let ([offending-stx (findf (λ (v) (not (identifier? v))) (syntax->list #'(maybe-id ...)))])
       (raise-syntax-error 
        #f "expected identifier to import" offending-stx))]
    [(_ . __)
     (raise-syntax-error #f "expected list of identifiers to import" stx)]
    [_ (raise-syntax-error #f "expected open parenthesis before import-primitive")]))

; Import a primitive procedure from the module scheme
(provide import-primitives)
(define-syntax (import-primitives stx)
  (syntax-case stx ()
    [(_ id ...) 
     (andmap identifier? (syntax->list #'(id ...)))
     (with-syntax ([(renamed-id ...) (generate-temporaries #'(id ...))]
                   [source (syntax-local-get-shadower
                            (syntax-local-introduce #'scheme))])
       #`(begin
           (require (only-in source [id renamed-id] ...))
           (do-not-annotate
            (define id
              (lambda args
                (unless (andmap (lambda (v) (and (location? v) (collector:flat? v))) args)
                  (error 'id (string-append "all arguments must be <heap-value?>s, "
                                            "even if the imported procedure accepts structured "
                                            "data")))
                (let ([result (apply renamed-id (map collector:deref args))])
                  (cond
                    [(void? result) (values)]
                    [(heap-value? result) (collector:alloc-flat result)]
                    [else 
                     (error 'id (string-append "imported primitive must return <heap-value?>, "
                                               "received ~a" result))])))))
           ...))]
    [(_ maybe-id ...) 
     (ormap (λ (v) (and (not (identifier? v)) v)) (syntax->list #'(maybe-id ...)))
     (let ([offending-stx (findf (λ (v) (not (identifier? v))) (syntax->list #'(maybe-id ...)))])
       (raise-syntax-error 
        #f "expected identifier to import" offending-stx))]
    [(_ . __)
     (raise-syntax-error #f "expected list of identifiers to import" stx)]
    [_ (raise-syntax-error #f "expected open parenthesis before import-primitive")]))

(define (mutator-lift f) 
  (lambda args
    ;(printf "Applying ~a~n" f)
    (let ([result (apply f (map collector:deref args))])
      ;(printf "Args: ~a, results: ~a~n" result (map collector:deref args))
      (if (void? result)
          (values)
          (collector:alloc-flat result)))))
(define-syntax (provide/lift stx)
  (syntax-case stx ()
    [(_ id ...)
     (andmap identifier? (syntax->list #'(id ...)))
     (with-syntax ([(lifted-id ...) (generate-temporaries #'(id ...))])
       #'(begin
           (define lifted-id (mutator-lift id)) ...
           (provide (rename-out [lifted-id id] ...))))]))

(provide/lift 
 symbol? boolean? number?
 add1 sub1 zero? + - * / even? odd? = < > <= >=)

(provide (rename-out (mutator-set-first! set-first!)))
(define (mutator-set-first! x y)
  (collector:set-first! x y)
  (values))

(provide (rename-out (mutator-set-rest! set-rest!)))
(define (mutator-set-rest! x y)
  (collector:set-rest! x y)
  (values))

(provide (rename-out [mutator-empty empty]))
(define-syntax mutator-empty
  (syntax-id-rules (mutator-empty)
    [_ (mutator-quote ())]))

(provide (rename-out (mutator-empty? empty?)))
(define (mutator-empty? loc)
  (cond
    [(collector:flat? loc) 
     (collector:alloc-flat (empty? (collector:deref loc)))]
    [else 
     (collector:alloc-flat false)]))

(provide (rename-out [mutator-cons? cons?]))
(define (mutator-cons? loc)
  (collector:alloc-flat (collector:cons? loc)))

(provide (rename-out [mutator-printf printf]))
(define-syntax (mutator-printf stx)
  (syntax-case stx ()
    [(_ fmt arg ...)
     ; We must invoke mutator-app to A-normalize the arguments.
     (syntax/loc stx 
       (begin
         (mutator-app printf (do-not-annotate (#%datum . fmt)) 
                      (mutator-app gc->scheme arg) ...)
         (values)))]))

(define-syntax mutator-quote
  (syntax-rules ()
    [(_ (a . d)) 
     (mutator-app collector:cons (mutator-quote a) (mutator-quote d))]
    [(_ s) 
     (collector:alloc-flat 's)]))

(define-syntax (mutator-datum stx)
  (syntax-case stx ()
    [(_ . e) 
     (quasisyntax/loc stx (#%app collector:alloc-flat (#%datum . e)))]))


(define-syntax (mutator-app stx)
  (syntax-case stx ()
    [(_ e ...)
     (let* ([no-expand (lambda (exp)
                         (syntax-case stx (do-not-annotate)
                           [(do-not-annotate _) #t]
                           [id  (and (identifier? exp)
                                     (not (eq? (syntax-e exp) 'empty)))]))]
            [exps (syntax->list #'(e ...))]
            [tmps (generate-temporaries #'(e ...))]
            [new-exps (map (lambda (exp tmp) (if (no-expand exp) exp tmp)) exps tmps)])
       (foldr (lambda (exp tmp acc)
                (if (no-expand exp)
                    acc
                    (with-syntax ([tmp tmp] [exp exp] [acc acc])
                      
                      (syntax/loc stx
                        (let ([tmp exp])
                          acc)))))
              (with-syntax ([(new-exp ...) new-exps])
                (syntax/loc stx (#%app new-exp ...)))
              exps tmps))]))


;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Annotation

(define (ensure-value loc)
  (if (or (location? loc) (procedure? loc))
      loc
      (error (format "expected <location?> or primitive procedure; received ~a" loc))))

(define-for-syntax expand-stop-list 
  (list #'#%top #'set! #'test/location=? #'do-not-annotate #'require 
        #'halt-on-errors #'print-only-errors))

(define-syntax (do-not-annotate stx)
  (syntax-case stx ()
    [(_ expr) #'expr]))

(define-for-syntax (remove-at lst remove-ix)
  (let loop ([ix 0] [lst lst])
    (cond
      [(empty? lst) empty]
      [(= ix remove-ix) (rest lst)]
      [else (cons (first lst) (loop (add1 ix) (rest lst)))])))

; annotate : syntax? (listof syntax?) boolean? -> syntax?
(define-for-syntax (annotate stx env-roots tail?)
  ;(printf "Env ~a for ~a~n" (map syntax-e env-roots) stx)
  (syntax-recertify
   (syntax-case stx (provide define-values require case-lambda if begin begin0 let-values
                             letrec-values set! quote #%app #%top #%datum #%plain-lambda
                             define-syntaxes define-values-for-syntax allocator-setup
                             #%expression test/location=? do-not-annotate)
     [(provide spec ...) stx]
     [(require spec ...) stx]
     [(allocator-setup . _) stx]
     [(test/location=? . _) stx]
     [(do-not-annotate expr) stx]
     [(define-values (id ...) expr)
      (quasisyntax/loc stx
        (begin
          (define-values (id ...) #,(annotate #'expr env-roots false))
          (add-global-root! (make-root 'id (λ () id) (λ (loc) (set! id loc))))
          ...))]
     [(#%plain-lambda (id ...) body ... body-tail)
      (let* ([free-ids (find-referenced-locals env-roots stx)]
             [body-env-roots (append (syntax->list #'(id ...)) env-roots)])
        (with-syntax ([(free-id ...) free-ids]
                      [(env-id ...) env-roots])
          #`(let* ([closure
                    (lambda (id ...)
                      #,@(map (λ (e) (annotate e body-env-roots false)) (syntax->list #'(body ...)))
                      #,(annotate #'body-tail body-env-roots true))])
              (add-closure-env! closure (list (make-env-root free-id) ...))
              #,(if tail?
                    #`(#%app collector:alloc-flat closure)
                    #`(with-continuation-mark gc-roots-key 
                        (list (make-root 'env-id (λ () env-id) (λ (loc) (set! env-id loc))) ...)
                        (#%app collector:alloc-flat closure))))))]
     [(case-lambda . _) (raise-syntax-error false "case-lambda is not yet supported" stx)]
     [(if test-expr expr ...)
      #`(if (collector:deref #,(annotate #'test-expr env-roots false))
            #,@(map (λ (e) (annotate e env-roots tail?)) (syntax->list #'(expr ...))))]
     [(begin) #'(values)]
     [(begin expr ... tail-expr)
      #`(begin
          #,@(map (λ (e) (annotate e env-roots false)) (syntax->list #'(expr ...)))
          #,(annotate #'tail-expr env-roots true))]
     [(begin0 . _) (raise-syntax-error false "begin0 is not yet supported" stx)]
     [(letrec-values . _)
      (raise-syntax-error false "letrec-values is not yet supported" stx)]
     [(let-values ([(id ...) bound-expr] ...) body-expr ... body-tail-expr)
      (andmap identifier? (syntax->list #'(id ... ...)))
      (with-syntax ([((tmp ...) ...) (map generate-temporaries (syntax->list #'((id ...) ...)))])
        (let* ([binding-list (syntax->list #'((tmp ...) ...))]
               [accumulated-local-roots 
                (build-list (length binding-list) 
                            (λ (n) (append-map syntax->list (take binding-list n))))]
               [env-roots-body (append 
                                (syntax->list #'(id ... ...))
                                ; not necessary (append-map syntax->list binding-list)
                                env-roots)])
          
          (with-syntax ([(annotated-bound-expr ...)
                         (map (λ (expr env-roots-ext) (annotate expr (append env-roots-ext env-roots) false))
                              (syntax->list #'(bound-expr ...))
                              accumulated-local-roots)])
            #`(let*-values ([(tmp ...) annotated-bound-expr] ...)
                (let-values ([(id ...) (values tmp ...)] ...)
                  #,@(map (λ (e) (annotate e env-roots-body tail?)) (syntax->list  #'(body-expr ...)))
                  #,(annotate #'body-tail-expr env-roots-body tail?))))))]
     [(set! id expr)
      #`(begin
          (set! id #,(annotate (local-expand #'expr 'expression expand-stop-list) env-roots tail?))
          (values))]
     [(quote stuff) #'(mutator-quote stuff)]
     [(#%top . exp) stx]
     [(#%datum . exp) #'exp]
     [(#%expression expr)
      (annotate #'expr env-roots tail?)]
     [(#%plain-app func-expr arg-expr ...) 
      ; All applications are in A-normal form.
      (begin
        (with-syntax ([(env-id ...) env-roots])
          (if tail?
              ; If this call is in tail position, we will not need access to its environment when it returns.
              (syntax/loc stx ((deref func-expr) arg-expr ...))
              ; If this call is not in tail position, we make the environment at the call site
              ; reachable.
              #`(with-continuation-mark gc-roots-key 
                  (list (make-root 'env-id (λ () env-id) (λ (loc) (set! env-id loc))) ...)
                  #,(syntax/loc stx ((deref func-expr) arg-expr ...))))))]
     [id (identifier? #'id) #'id]
     [expr (raise-syntax-error #f (format "Bug in the PLAI Mutator framework; unrecognized form ~a" (syntax->datum #'expr))
                               #'expr)])
   ; Remaining arguments from syntax-recertify
   stx 
   (current-code-inspector)
   false))

(define (deref proc/loc)
  (cond
    [(procedure? proc/loc) proc/loc]
    [(location? proc/loc) (collector:deref proc/loc)]
    [else (error 'deref "expected <location?> or <procedure?; received ~a" proc/loc)]))


(define (is-non-primitive? x) x)

(define (scheme->gc v)
  (cond
    [(heap-value? v) (collector:alloc-flat v)]
    [(cons? v) (collector:cons (scheme->gc (car v)) (scheme->gc (cdr v)))]
    [(mpair? v) (collector:cons (scheme->gc (mcar v)) (scheme->gc (mcdr v)))]
    [else (error (format "unsupported data type"))]))

(define (gc->scheme location)
  (local ([define no-value (gensym)]
          [define loc->val (make-hash)]
          [define (unwrap loc)
            (let ([maybe-val (hash-ref loc->val loc no-value)])
              (if (eq? maybe-val no-value)
                  (cond
                    [(collector:flat? loc)
                     (let ([val (collector:deref loc)])
                       (hash-set! loc->val loc val)
                       val)]
                    [(collector:cons? loc)
                     (let ([val (mcons 'uninitalized 'unintialized)])
                       (hash-set! loc->val loc val)
                       (set-mcar! val (unwrap (collector:first loc)))
                       (set-mcdr! val (unwrap (collector:rest loc)))
                       val)]
                    [else 
                     (error (format "gc:flat? and gc:cons? both returned false for ~a" loc))])
                  maybe-val))])
    (unwrap location)))

;;; Special require for garbage collector
(define-for-syntax (allocator-setup-internal stx)
  (with-syntax ([(init-allocator gc:deref gc:alloc-flat gc:cons 
                                 gc:first gc:rest 
                                 gc:flat? gc:cons?
                                 gc:set-first! gc:set-rest!)
                 (map (λ (s) (datum->syntax stx s))
                      '(init-allocator gc:deref gc:alloc-flat gc:cons 
                                       gc:first gc:rest 
                                       gc:flat? gc:cons?
                                       gc:set-first! gc:set-rest!))])
    (syntax-case stx ()
      [(collector-module heap-size)
       (begin
         (set! required-allocator-stx 
               (if (alternate-collector)
                   (datum->syntax stx (alternate-collector))
                   #'collector-module))
         #`(do-not-annotate
            (begin
              #,(if (alternate-collector)
                    #`(require #,(datum->syntax #'collector-module (alternate-collector)))
                    #`(require collector-module))
              
              (set-collector:deref! gc:deref)
              (set-collector:alloc-flat! gc:alloc-flat)
              (set-collector:cons! gc:cons)
              (set-collector:first! gc:first)
              (set-collector:rest! gc:rest)
              (set-collector:flat?! gc:flat?)
              (set-collector:cons?! gc:cons?)
              (set-collector:set-first!! gc:set-first!)
              (set-collector:set-rest!! gc:set-rest!)
              
              (init-heap! (#%datum . heap-size))
              (when (gui-available?) 
                (if (<= (#%datum . heap-size) 200)
                    (set-ui! (dynamic-require `(planet #,(this-package-version-symbol private/gc-gui)) 'heap-viz%))
                    (printf "Large heap; the heap visualizer will not be displayed.~n")))
              (init-allocator)           
              )))]
      [_ (raise-syntax-error false "expected (allocator-setup <collector-file> <heap-size>)" stx)])))


(define-syntax (mutator-top-interaction stx)
  (syntax-case stx (require provide define)
    [(_ . (require spec ...)) #'(require spec ...)]
    [(_ . (provide spec ...)) #'(provide spec ...)]
    [(_ . (define id expr))
     #`(define id
         #,(annotate (local-expand #'expr 'top-level expand-stop-list) empty false))]
    [(_ . expr)
     #`(let ([result-addr #,(annotate (local-expand #'expr 'top-level expand-stop-list) empty false)])
         (if (procedure? result-addr)
             (begin
               (printf "Imported procedure~n")
               result-addr)
             (begin
               (printf "Value at location ~a:~n" result-addr)
               (gc->scheme result-addr))))]))


(define-syntax (mutator-module-begin stx)
  (syntax-case stx (allocator-setup)
    [(_ (allocator-setup . setup) module-expr ...)
     #`(#%module-begin
        #,(allocator-setup-internal #'setup)
        #,@(map (λ (e)
                  (let* ([expanded-expr (local-expand e 'top-level expand-stop-list)]
                         [annotated (annotate expanded-expr empty false)])
                    annotated))
                (syntax->list #'(module-expr ...)))
        
        )]))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Testing support

(define-syntax (test/location=? stx)
  (syntax-case stx ()
    [(_ e1 e2)
     (let ([e1-expanded (annotate (local-expand #'e1 'top-level expand-stop-list) empty false)]
           [e2-expanded (annotate (local-expand #'e2 'top-level expand-stop-list) empty false)])
       #`(let ([e1-addr #,e1-expanded]
               [e2-addr #,e2-expanded])
           (generic-test (λ () (eq? e1-addr e2-addr)) (λ (v) v) (quote #,(syntax->datum #'e1))
                         (quote #,(syntax->datum #'e2)) (format "at line ~a" #,(syntax-line stx)))))]))

(define-for-syntax (flat-heap-value? v)
  (or (number? v) (boolean? v)))

(define-for-syntax (expand-scheme stx)
  (syntax-case stx (#%datum)
    [val (flat-heap-value? (syntax->datum #'val)) #'(#%datum . val)]
    [(quote e) #'(scheme:quote e)]
    [_ (raise-syntax-error 'test/value=? "must be a number, boolean or a quoted value" stx)]))

; This is because set-car! and set-cdr! have been removed.  Only v1 may have mcons.
(define (equal/mcons? v1 v2)
  (cond
    [(mpair? v1)
     (if (cons? v2)
         (and (equal/mcons? (mcar v1) (car v2)) (equal/mcons? (mcdr v1) (cdr v2)))
         false)]
    [(cons? v1) ; This will not happen, but for completeness
     (if (cons? v2)
         (and (equal/mcons? (car v1) (car v2)) (equal/mcons? (cdr v1) (cdr v2)))
         false)]
    [else (equal? v1 v2)]))

(define-syntax (test/value=? stx)
  (syntax-case stx ()
    [(_ mutator-expr scheme-datum)
     (quasisyntax/loc stx
       (let ([v1 mutator-expr]
             [v2 (do-not-annotate #,(expand-scheme #'scheme-datum))])
         ; A-normalize
         (mutator-app generic-test 
                      (do-not-annotate (λ () (equal/mcons? (gc->scheme v1) v2)))
                      (do-not-annotate (λ (v) v))
                      (do-not-annotate (quote #,(syntax->datum #'mutator-expr)))
                      (do-not-annotate (quote #,(syntax->datum #'scheme-datum)))
                      (do-not-annotate #,(syntax-line stx)))
         (values)))]))