(.module:
  lux
  (lux (control ["p" parser])
       (data [maybe]
             ["e" error]
             [number]
             [product]
             text/format
             (coll [list "list/" Functor<List> Fold<List> Monoid<List>]
                   [dict #+ Dict]))
       (macro [code]
              ["s" syntax]))
  (luxc (lang ["la" analysis]
              ["ls" synthesis]
              [".L" extension #+ Syntheses]
              (synthesis [".S" case]
                         [".S" function]
                         [".S" loop])
              [".L" variable #+ Variable])
        ))

(def: init-env (List Variable) (list))
(def: init-resolver (Dict Int Int) (dict.new number.Hash<Int>))

(def: (prepare-body inner-arity arity body)
  (-> ls.Arity ls.Arity ls.Synthesis ls.Synthesis)
  (if (functionS.nested? inner-arity)
    body
    (loopS.reify-recursion arity body)))

(def: (let$ register inputS bodyS)
  (-> Nat ls.Synthesis ls.Synthesis ls.Synthesis)
  (` ("lux let" (~ (code.nat register)) (~ inputS) (~ bodyS))))

(def: (if$ testS thenS elseS)
  (-> ls.Synthesis ls.Synthesis ls.Synthesis ls.Synthesis)
  (` ("lux if" (~ testS)
      (~ thenS)
      (~ elseS))))

(def: (function$ arity environment body)
  (-> ls.Arity (List Variable) ls.Synthesis ls.Synthesis)
  (` ("lux function" (~ (code.nat arity))
      [(~+ (list/map code.int environment))]
      (~ body))))

(def: (variant$ tag last? valueS)
  (-> Nat Bool ls.Synthesis ls.Synthesis)
  (` ((~ (code.nat tag)) (~ (code.bool last?)) (~ valueS))))

(def: (var$ var)
  (-> Variable ls.Synthesis)
  (` ((~ (code.int var)))))

(def: (procedure$ name argsS)
  (-> Text (List ls.Synthesis) ls.Synthesis)
  (` ((~ (code.text name)) (~+ argsS))))

(def: (call$ funcS argsS)
  (-> ls.Synthesis (List ls.Synthesis) ls.Synthesis)
  (` ("lux call" (~ funcS) (~+ argsS))))

(def: (synthesize-case arity num-locals synthesize inputA branchesA)
  (-> ls.Arity Nat (-> Nat la.Analysis ls.Synthesis)
      la.Analysis (List [la.Pattern la.Analysis])
      ls.Synthesis)
  (let [inputS (synthesize num-locals inputA)]
    (case (list.reverse branchesA)
      (^multi (^ (list [(^code ("lux case bind" (~ [_ (#.Nat input-register)])))
                        (^code ((~ [_ (#.Int var)])))]))
              (not (variableL.captured? var))
              (n/= input-register (variableL.local-register var)))
      inputS

      (^ (list [(^code ("lux case bind" (~ [_ (#.Nat register)]))) bodyA]))
      (let$ (if (functionS.nested? arity)
              (n/+ (n/dec arity) register)
              register)
            inputS
            (synthesize (n/inc num-locals) bodyA))
      
      (^or (^ (list [(^code true) thenA] [(^code false) elseA]))
           (^ (list [(^code false) elseA] [(^code true) thenA])))
      (if$ inputS (synthesize num-locals thenA) (synthesize num-locals elseA))

      (#.Cons [lastP lastA] prevsPA)
      (let [transform-branch (: (-> la.Pattern la.Analysis ls.Path)
                                (caseS.path arity num-locals synthesize))
            pathS (list/fold caseS.weave
                             (transform-branch lastP lastA)
                             (list/map (product.uncurry transform-branch) prevsPA))]
        (` ("lux case" (~ inputS) (~ pathS))))

      _
      (undefined)
      )))

(def: (synthesize-apply synthesize num-locals exprA)
  (-> (-> la.Analysis ls.Synthesis) Nat la.Analysis ls.Synthesis)
  (let [[funcA argsA] (functionS.unfold-apply exprA)
        funcS (synthesize funcA)
        argsS (list/map synthesize argsA)]
    (case funcS
      (^multi (^code ("lux function" (~ [_ (#.Nat _arity)]) [(~+ _env)] (~ _bodyS)))
              (and (n/= _arity (list.size argsS))
                   (not (loopS.contains-self-reference? _bodyS)))
              [(s.run _env (p.some s.int)) (#e.Success _env)])
      (` ("lux loop" (~ (code.nat (n/inc num-locals))) [(~+ argsS)]
          (~ (loopS.adjust _env num-locals _bodyS))))

      (^code ("lux call" (~ funcS') (~+ argsS')))
      (call$ funcS' (list/compose argsS' argsS))

      _
      (call$ funcS argsS))))

(def: #export (synthesize extensions expressionA)
  (-> Syntheses la.Analysis ls.Synthesis)
  (loop [arity +0
         resolver init-resolver
         direct? false
         num-locals +0
         expressionA expressionA]
    (case expressionA
      (^code [(~ _left) (~ _right)])
      (` [(~+ (list/map (recur arity resolver false num-locals)
                        (la.unfold-tuple expressionA)))])

      (^or (^code ("lux sum left" (~ _)))
           (^code ("lux sum right" (~ _))))
      (let [[tag last? value] (maybe.assume (la.unfold-variant expressionA))]
        (variant$ tag last? (recur arity resolver false num-locals value)))

      (^code ((~ [_ (#.Int var)])))
      (if (variableL.local? var)
        (if (functionS.nested? arity)
          (if (variableL.self? var)
            (call$ (var$ 0) (|> (list.n/range +1 (n/dec arity))
                                (list/map (|>> variableL.local code.int (~) () (`)))))
            (var$ (functionS.adjust-var arity var)))
          (var$ var))
        (var$ (maybe.default var (dict.get var resolver))))

      (^code ("lux case" (~ inputA) (~ [_ (#.Record branchesA)])))
      (synthesize-case arity num-locals (recur arity resolver false) inputA branchesA)

      (^multi (^code ("lux function" [(~+ scope)] (~ bodyA)))
              [(s.run scope (p.some s.int)) (#e.Success raw-env)])
      (let [function-arity (if direct?
                             (n/inc arity)
                             +1)
            env (list/map (function [closure]
                            (case (dict.get closure resolver)
                              (#.Some resolved)
                              (if (and (variableL.local? resolved)
                                       (functionS.nested? arity)
                                       (|> resolved variableL.local-register (n/>= arity)))
                                (functionS.adjust-var arity resolved)
                                resolved)

                              #.None
                              (if (and (variableL.local? closure)
                                       (functionS.nested? arity))
                                (functionS.adjust-var arity closure)
                                closure)))
                          raw-env)
            env-vars (: (List Variable)
                        (case raw-env
                          #.Nil (list)
                          _ (|> (list.size raw-env) n/dec (list.n/range +0) (list/map variableL.captured))))
            resolver' (if (and (functionS.nested? function-arity)
                               direct?)
                        (list/fold (function [[from to] resolver']
                                     (dict.put from to resolver'))
                                   init-resolver
                                   (list.zip2 env-vars env))
                        (list/fold (function [var resolver']
                                     (dict.put var var resolver'))
                                   init-resolver
                                   env-vars))]
        (case (recur function-arity resolver' true function-arity bodyA)
          (^ [_ (#.Form (list [_ (#.Text "lux function")] [_ (#.Nat unmerged-arity)] env' bodyS'))])
          (let [merged-arity (n/inc unmerged-arity)]
            (function$ merged-arity env
                       (prepare-body function-arity merged-arity bodyS')))

          bodyS
          (function$ +1 env (prepare-body function-arity +1 bodyS))))

      (^code ("lux apply" (~+ _)))
      (synthesize-apply (recur arity resolver false num-locals) num-locals expressionA)

      (^code ((~ [_ (#.Text name)]) (~+ args)))
      (case (dict.get name extensions)
        #.None
        (procedure$ name (list/map (recur arity resolver false num-locals) args))
        
        (#.Some extension)
        (extension (recur arity resolver false num-locals) args))

      _
      expressionA)))