(module hoist mzscheme
  (require (planet "contract-utils.ss" ("cobbe" "contract-utils.plt" 1 0))
           (all-except (planet "list.ss" ("dherman" "list.plt" 1 0)) any)
           (lib "contract.ss")
           (lib "match.ss")
           (lib "struct.ss")
           "../syntax/ast.ss"
           "../syntax/token.ss")

  (define-struct (FunctionDeclaration/hoisted FunctionDeclaration) (functions variables))
  (define-struct (FunctionExpression/hoisted FunctionExpression) (functions variables))
  (define-struct (BlockStatement/hoisted BlockStatement) (functions variables))
;  (define-struct (LetExpression/hoisted LetExpression) (functions variables))
;  (define-struct (LetStatement/hoisted LetStatement) (functions variables))

  ;; (listof FunctionDeclaration/hoisted) * (listof Identifier) * (listof FunctionDeclaration/hoisted) * (listof Identifier)
  (define-struct hoisting (functions vars let-functions let-vars) #f)

  ;; null-hoisting : hoisting
  (define null-hoisting (make-hoisting null null null null))

  ;; append-hoisting : hoisting * hoisting -> hoisting
  (define (append-hoisting . hs)
    (match hs
      [(($ hoisting funss varss let-funss let-varss) ...)
       (make-hoisting (apply append funss)
                      (apply append varss)
                      (apply append let-funss)
                      (apply append let-varss))]))

  ;; cons-function-hoisting : FunctionDeclaration/hoisted * hoisting -> hoisting
  (define (cons-function-hoisting fun h)
    (copy-struct hoisting h [hoisting-functions (cons fun (hoisting-functions h))]))

  ;; cons-var-hoisting : Identifier * hoisting -> hoisting
  (define (cons-var-hoisting var h)
    (copy-struct hoisting h [hoisting-vars (cons var (hoisting-vars h))]))

  ;; cons-let-function-hoisting : FunctionDeclaration/hoisted * hoisting -> hoisting
  (define (cons-let-function-hoisting fun h)
    (copy-struct hoisting h [hoisting-let-functions (cons fun (hoisting-let-functions h))]))

  ;; cons-let-var-hoisting : Identifier * hoisting -> hoisting
  (define (cons-let-var-hoisting var h)
    (copy-struct hoisting h [hoisting-let-vars (cons var (hoisting-let-vars h))]))

  ;; a (continuation a -> b) is:
  ;;
  ;;   - (hoisting a -> b)

  (define (continuation/c c)
    (hoisting? c . -> . any))

  ;; an (a ->k b) is:
  ;;
  ;;   - (a (continuation b -> c) -> c)

  (define-syntax provide/contract/k
    (syntax-rules (->k)
      [(_ [name (domain . ->k . range)] ...)
       (provide/contract
         [name (domain (continuation/c range) . -> . any)]
         ...)]))

  ;; TODO: just directly use lset-difference instead of unique-vars at all
  ;; unique-vars : (listof FunctionDeclaration/hoisted) * (listof Identifier) ... -> (listof Identifier)
  (define (unique-vars funs . varss)
    (lset-difference Identifier=?
      (delete-duplicates (apply append varss) Identifier=?)
      (map FunctionDeclaration-name funs)))

  ;; ===========================================================================
  ;; TOP-LEVEL HOISTING FUNCTIONS
  ;; ===========================================================================

  ;; hoist-function-declaration : FunctionDeclaration -> FunctionDeclaration/hoisted
  (define (hoist-function-declaration decl)
    (match decl
      [($ FunctionDeclaration location name args body)
       (hoist-source-elements body
         (lambda (h stmts)
           (match h
             [($ hoisting funs vars let-funs let-vars)
              ;; TODO: add in the let-funs and let-vars as well
              (make-FunctionDeclaration/hoisted location name args stmts funs (unique-vars funs vars let-vars))])))]))

  ;; hoist-script : (listof SourceElement) -> (listof FunctionDeclaration/hoisted)
  ;;                                          (listof Identifier)
  ;;                                          (listof Statement)
  (define (hoist-script elts)
    (hoist-source-elements elts
      (lambda (h stmts)
        (match h
          [($ hoisting funs vars let-funs let-vars)
           ;; TODO: add in the let-funs and let-vars as well
           (values funs (unique-vars funs vars let-vars) stmts)]))))

  ;; ===========================================================================
  ;; COMPOUND HOISTING FUNCTIONS
  ;; ===========================================================================

  ;; map-k : (a ->k b) (listof a) ->k (listof b)
  (define (map-k hoist1 elts k)
    (if (null? elts)
        (k null-hoisting null)
        (hoist1 (car elts)
          (lambda (hoisting1 result)
            (map-k hoist1 (cdr elts)
              (lambda (hoisting2 results)
                (k (append-hoisting hoisting1 hoisting2)
                   (cons result results))))))))

  ;; filter-map-k : (a ->k (optional b)) (listof a) ->k (listof b)
  (define (filter-map-k hoist1 elts k)
    (if (null? elts)
        (k null-hoisting null)
        (hoist1 (car elts)
          (lambda (hoisting1 result)
            (filter-map-k hoist1 (cdr elts)
              (lambda (hoisting2 results)
                (k (append-hoisting hoisting1 hoisting2)
                   (if result (cons result results) results))))))))

  ;; append-map-k : (a ->k (listof b)) (listof a) ->k (listof b)
  (define (append-map-k hoist1 elts k)
    (if (null? elts)
        (k null-hoisting null)
        (hoist1 (car elts)
          (lambda (hoisting1 results1)
            (append-map-k hoist1 (cdr elts)
              (lambda (hoisting2 results2)
                (k (append-hoisting hoisting1 hoisting2)
                   (append results1 results2))))))))

  ;; hoist-source-elements : (listof SourceElement) ->k (listof Statement)
  (define (hoist-source-elements elts k)
    (map-k hoist-source-element elts k))

  ;; hoist-expressions : (listof Expression) ->k (listof Expression)
  (define (hoist-expressions exprs k)
    (map-k hoist-expression exprs k))

  ;; hoist-optional-expression : (optional Expression) ->k (optional Expression)
  (define (hoist-optional-expression expr k)
    (if (not expr)
        (k null-hoisting #f)
        (hoist-expression expr k)))

  ;; hoist-optional-expressions : (listof (optional Expression)) ->k (listof (optional Expression))
  (define (hoist-optional-expressions exprs k)
    (map-k hoist-optional-expression exprs k))

  ;; hoist-substatements : (listof SourceElement) ->k (listof Statement)
  (define (hoist-substatements stmts k)
    (filter-map-k hoist-substatement stmts k))

  ;; hoist-variable-bindings : (listof VariableInitializer) ->k (listof Expression)
  (define (hoist-variable-bindings decls k)
    (filter-map-k hoist-variable-binding decls k))

  ;; hoist-case-clauses : (listof CaseClause) ->k (listof CaseClause)
  (define (hoist-case-clauses cases k)
    (map-k hoist-case-clause cases k))

  ;; hoist-catch-clauses : (listof CatchClause) ->k (listof CatchClause)
  (define (hoist-catch-clauses catches k)
    (map-k hoist-catch-clause catches k))

  ;; ===========================================================================
  ;; CORE HOISTING FUNCTIONS
  ;; ===========================================================================

  ;; hoist-source-element : SourceElement ->k Statement
  (define (hoist-source-element src k)
    (if (Declaration? src)
        (hoist-declaration src k)
        (hoist-statement src k)))

  ;; hoist-declaration : Declaration ->k Statement
  (define (hoist-declaration decl0 k)
    (cond
      [(FunctionDeclaration? decl0)
       (k (cons-function-hoisting (hoist-function-declaration decl0) null-hoisting)
          (make-EmptyStatement (Term-location decl0)))]
      [(VariableDeclaration? decl0)
       (hoist-variable-declaration decl0
         (lambda (h expr)
           (k h (make-ExpressionStatement (Term-location decl0) expr))))]))

  ;; hoist-variable-declaration : VariableDeclaration ->k Expression
  (define (hoist-variable-declaration decl0 k)
    (match decl0
      [($ VariableDeclaration loc bindings)
       (hoist-variable-bindings bindings
         (lambda (h exprs)
           (k h (if (and (pair? exprs)
                         (null? (cdr exprs)))
                    (car exprs)
                    (make-ListExpression loc exprs)))))]))

  ;; hoist-substatement : SourceElement ->k (optional Statement)
  (define (hoist-substatement src0 k)
    (if (FunctionDeclaration? src0)
        (k (cons-function-hoisting (hoist-function-declaration src0) null-hoisting) #f)
        (hoist-statement src0 k)))

  ;; hoist-variable-binding : VariableInitializer ->k (optional Expression)
  (define (hoist-variable-binding decl k)
    (match decl
      [($ VariableInitializer loc id #f)
       (k (cons-var-hoisting id null-hoisting) #f)]
      [($ VariableInitializer loc id init)
       (hoist-expression init
         (lambda (h init)
           (k (cons-var-hoisting id h)
              (make-AssignmentExpression loc (make-VarReference (Term-location id) id) '= init))))]))

  ;; hoist-case-clause : CaseClause ->k CaseClause
  (define (hoist-case-clause case k)
    (match case
      [($ CaseClause loc #f answer)
       (hoist-substatements answer
         (lambda (h answer)
           (k h (make-CaseClause loc #f answer))))]
      [($ CaseClause loc question answer)
       (hoist-expression question
         (lambda (h1 question)
           (hoist-substatements answer
             (lambda (h2 answer)
               (k (append-hoisting h1 h2)
                  (make-CaseClause loc question answer))))))]))

  ;; hoist-catch-clause : CatchClause ->k CatchClause
  (define (hoist-catch-clause catch k)
    (match catch
      [($ CatchClause loc id body0)
       (hoist-statement body0
         (lambda (h body)
           (k h (make-CatchClause loc id body))))]))

  (define (to-location x)
    (cond
      [(not x) #f]
      [(position? x) x]
      [else (ast-location x)]))

  ;; optional-statement->statement : (optional Statement) (optional has-location) -> Statement
  (define (optional-statement->statement stmt loc)
    (or stmt (make-EmptyStatement (to-location loc))))

  ;; hoist-statement : Statement ->k Statement
  (define (hoist-statement stmt k)
    (match stmt
      [($ BlockStatement loc stmts)
       (hoist-substatements stmts
         (lambda (h stmts)
           (k h (make-BlockStatement loc stmts))))]
      [($ ExpressionStatement loc expr)
       (hoist-expression expr
         (lambda (h expr)
           (k h (make-ExpressionStatement loc expr))))]
      [($ IfStatement loc test consequent0 alternate0)
       (hoist-expression test
         (lambda (h1 test)
           (hoist-substatement consequent0
             (lambda (h2 consequent)
               (if alternate0
                   (hoist-substatement alternate0
                     (lambda (h3 alternate)
                       (k (append-hoisting h1 h2 h3)
                          (make-IfStatement loc
                                            test
                                            (optional-statement->statement consequent consequent0)
                                            alternate))))
                   (k (append-hoisting h1 h2)
                      (make-IfStatement loc test (optional-statement->statement consequent consequent0) #f)))))))]
      [($ DoWhileStatement loc body0 test)
       (hoist-substatement body0
         (lambda (h1 body)
           (hoist-expression test
             (lambda (h2 test)
               (k (append-hoisting h1 h2)
                  (make-DoWhileStatement loc (optional-statement->statement body body0) test))))))]
      [($ WhileStatement loc test body0)
       (hoist-expression test
         (lambda (h1 test)
           (hoist-substatement body0
             (lambda (h2 body)
               (k (append-hoisting h1 h2)
                  (make-WhileStatement loc test (optional-statement->statement body body0)))))))]
      [($ ForStatement loc init test incr body0)
       (let ([hoist (if (VariableDeclaration? init)
                        hoist-variable-declaration
                        hoist-optional-expression)])
         (hoist init
           (lambda (h1 init)
             (hoist-optional-expression test
               (lambda (h2 test)
                 (hoist-optional-expression incr
                   (lambda (h3 incr)
                     (hoist-substatement body0
                       (lambda (h4 body)
                         (k (append-hoisting h1 h2 h3 h4)
                            (make-ForStatement loc init test incr (optional-statement->statement body body0))))))))))))]
      [($ ForInStatement loc (and lhs (? Expression?)) container body0)
       (hoist-expression lhs
         (lambda (h1 lhs)
           (hoist-expression container
             (lambda (h2 container)
               (hoist-substatement body0
                 (lambda (h3 body)
                   (k (append-hoisting h1 h2 h3)
                      (make-ForInStatement loc lhs container (optional-statement->statement body body0)))))))))]
      [($ ForInStatement loc ($ VariableDeclaration _ ($ VariableInitializer v-loc id #f)) container body0)
       (hoist-expression container
         (lambda (h1 container)
           (hoist-substatement body0
             (lambda (h2 body)
               (k (cons-var-hoisting id (append-hoisting h1 h2))
                  (make-ForInStatement loc
                                       (make-VarReference v-loc id)
                                       container
                                       (optional-statement->statement body body0)))))))]
      [($ ReturnStatement loc (and expr (? Expression?)))
       (hoist-expression expr
         (lambda (h expr)
           (k h (make-ReturnStatement loc expr))))]
      [($ WithStatement loc test body0)
       (hoist-expression test
         (lambda (h1 test)
           (hoist-substatement body0
             (lambda (h2 body)
               (k (append-hoisting h1 h2)
                  (make-WithStatement loc test (optional-statement->statement body body0)))))))]
      [($ SwitchStatement loc expr cases)
       (hoist-expression expr
         (lambda (h1 expr)
           (hoist-case-clauses cases
             (lambda (h2 cases)
               (k (append-hoisting h1 h2)
                  (make-SwitchStatement loc expr cases))))))]
      [($ LabelledStatement loc label stmt0)
       (hoist-substatement stmt0
         (lambda (h stmt)
           (k h (make-LabelledStatement loc label (optional-statement->statement stmt stmt0)))))]
      [($ ThrowStatement loc expr)
       (hoist-expression expr
         (lambda (h expr)
           (k h (make-ThrowStatement loc expr))))]
      [($ TryStatement loc body0 catch0 finally0)
       (hoist-statement body0
         (lambda (h1 body)
           (hoist-catch-clauses catch0
             (lambda (h2 catch)
               (if finally0
                   (hoist-statement finally0
                     (lambda (h3 finally)
                       (k (append-hoisting h1 h2 h3)
                          (make-TryStatement loc body catch finally))))
                   (k (append-hoisting h1 h2)
                      (make-TryStatement loc body catch #f)))))))]
      [_ (k null-hoisting stmt)]))

  ;; hoist-expression : Expression ->k Expression
  (define (hoist-expression expr k)
    (match expr
      [($ ArrayLiteral loc elts)
       (hoist-optional-expressions elts
         (lambda (h elts)
           (k h (make-ArrayLiteral loc elts))))]
      [($ ObjectLiteral loc ([props . vals] ...))
       (hoist-expressions vals
         (lambda (h vals)
           (k h (make-ObjectLiteral loc (map cons props vals)))))]
      [($ BracketReference loc container key)
       (hoist-expression container
         (lambda (h1 container)
           (hoist-expression key
             (lambda (h2 key)
               (k (append-hoisting h1 h2)
                  (make-BracketReference loc container key))))))]
      [($ DotReference loc container id)
       (hoist-expression container
         (lambda (h container)
           (k h (make-DotReference loc container id))))]
      [($ NewExpression loc constructor args)
       (hoist-expression constructor
         (lambda (h1 constructor)
           (hoist-expressions args
             (lambda (h2 args)
               (k (append-hoisting h1 h2)
                  (make-NewExpression loc constructor args))))))]
      [($ PostfixExpression loc expr op)
       (hoist-expression expr
         (lambda (h expr)
           (k h (make-PostfixExpression loc expr op))))]
      [($ PrefixExpression loc op expr)
       (hoist-expression expr
         (lambda (h expr)
           (k h (make-PrefixExpression loc op expr))))]
      [($ InfixExpression loc left op right)
       (hoist-expression left
         (lambda (h1 left)
           (hoist-expression right
             (lambda (h2 right)
               (k (append-hoisting h1 h2)
                  (make-InfixExpression loc left op right))))))]
      [($ ConditionalExpression loc test consequent alternate)
       (hoist-expression test
         (lambda (h1 test)
           (hoist-expression consequent
             (lambda (h2 consequent)
               (hoist-expression alternate
                 (lambda (h3 alternate)
                   (k (append-hoisting h1 h2 h3)
                      (make-ConditionalExpression loc test consequent alternate))))))))]
      [($ AssignmentExpression loc left op right)
       (hoist-expression left
         (lambda (h1 left)
           (hoist-expression right
             (lambda (h2 right)
               (k (append-hoisting h1 h2)
                  (make-AssignmentExpression loc left op right))))))]
      [($ FunctionExpression loc name args body)
       (k null-hoisting (hoist-source-elements body
                          (lambda (h body)
                            (match h
                              [($ hoisting funs vars let-funs let-vars)
                               ;; TODO: use the let-funs and let-vars
                               (make-FunctionExpression/hoisted loc name args body funs (unique-vars funs vars let-vars))]))))]
;      [($ LetExpression loc bindings body)
;       (k null null (hoist-source-element body
;                      (lambda (funs vars body)
;                        (make-LetExpression/hoisted loc bindings body funs (unique-vars funs vars)))))]
      [($ CallExpression loc method args)
       (hoist-expression method
         (lambda (h1 method)
           (hoist-expressions args
             (lambda (h2 args)
               (k (append-hoisting h1 h2)
                  (make-CallExpression loc method args))))))]
      [($ ParenExpression loc expr)
       (hoist-expression expr
         (lambda (h expr)
           (k h (make-ParenExpression loc expr))))]
      [($ ListExpression loc exprs)
       (hoist-expressions exprs
         (lambda (h exprs)
           (k h (make-ListExpression loc exprs))))]
      [_ (k null-hoisting expr)]))

  (provide/contract
    [continuation/c ((union flat-contract? predicate/c) . -> . contract?)])

  (provide/contract/k
    [hoist-source-elements ((listof SourceElement?) . ->k . (listof Statement?))]
    [hoist-declaration (Declaration? . ->k . Statement?)]
    [hoist-variable-declaration (VariableDeclaration? . ->k . Expression?)]
    [hoist-expressions ((listof Expression?) . ->k . (listof Expression?))]
    [hoist-optional-expressions ((listof (optional/c Expression?)) . ->k . (listof (optional/c Expression?)))]
    [hoist-source-element (SourceElement? . ->k . Statement?)]
    [hoist-statement (Statement? . ->k . Statement?)]
    [hoist-expression (Expression? . ->k . Expression?)]
    [hoist-substatement (SubStatement? . ->k . (optional/c Statement?))]
    [hoist-substatements ((listof SubStatement?) . ->k . (listof Statement?))]
    [hoist-variable-binding (VariableInitializer? . ->k . (optional/c Expression?))]
    [hoist-variable-bindings ((listof VariableInitializer?) . ->k . (listof Expression?))]
    [hoist-case-clause (CaseClause? . ->k . CaseClause?)]
    [hoist-case-clauses ((listof CaseClause?) . ->k . (listof CaseClause?))]
    [hoist-catch-clause (CatchClause? . ->k . CatchClause?)]
    [hoist-catch-clauses ((listof CatchClause?) . ->k . (listof CatchClause?))])

  (provide/contract
    [hoist-function-declaration (FunctionDeclaration? . -> . FunctionDeclaration/hoisted?)]
    [hoist-script ((listof SourceElement?) . -> . (values (listof FunctionDeclaration/hoisted?)
                                                          (listof Identifier?)
                                                          (listof Statement?)))])

  (provide/contract
    (struct (FunctionDeclaration/hoisted FunctionDeclaration) ([location (optional/c region?)]
                                                               [name Identifier?]
                                                               [args (listof Identifier?)]
                                                               [body (listof Statement?)]
                                                               [functions (listof FunctionDeclaration/hoisted?)]
                                                               [vars (listof Identifier?)]))
    (struct (FunctionExpression/hoisted FunctionExpression) ([location (optional/c region?)]
                                                             [name (optional/c Identifier?)]
                                                             [args (listof Identifier?)]
                                                             [body (listof Statement?)]
                                                             [functions (listof FunctionDeclaration/hoisted?)]
                                                             [vars (listof Identifier?)]))
    #;(struct (LetExpression/hoisted LetExpression) ([location (optional/c region?)]
                                                   [bindings (listof VariableDeclaration?)]
                                                   [body Statement?]
                                                   [functions (listof FunctionDeclaration/hoisted?)]
                                                   [vars (listof Identifier?)]))
    ))