(.module:
  lux
  (lux (control [monad #+ do]
                ["p" parser])
       (data [maybe]
             (coll [list "list/" Functor<List>]))
       (macro [code]
              [syntax]))
  (luxc (lang ["ls" synthesis]
              [".L" variable #+ Variable Register])))

(def: #export (contains-self-reference? exprS)
  (-> ls.Synthesis Bool)
  (case exprS
    (^ [_ (#.Form (list [_ (#.Nat tag)] [_ (#.Bool last?)] memberS))])
    (contains-self-reference? memberS)

    [_ (#.Tuple membersS)]
    (list.any? contains-self-reference? membersS)

    (^ [_ (#.Form (list [_ (#.Int var)]))])
    (variableL.self? var)

    (^ [_ (#.Form (list [_ (#.Text "lux case")] inputS pathS))])
    (or (contains-self-reference? inputS)
        (loop [pathS pathS]
          (case pathS
            (^or (^ [_ (#.Form (list [_ (#.Text "lux case alt")] leftS rightS))])
                 (^ [_ (#.Form (list [_ (#.Text "lux case seq")] leftS rightS))]))
            (or (recur leftS)
                (recur rightS))

            (^ [_ (#.Form (list [_ (#.Text "lux case exec")] bodyS))])
            (contains-self-reference? bodyS)
            
            _
            false)))

    (^ [_ (#.Form (list [_ (#.Text "lux function")] arity [_ (#.Tuple environment)] bodyS))])
    (list.any? (function (_ captured)
                 (case captured
                   (^ [_ (#.Form (list [_ (#.Int var)]))])
                   (variableL.self? var)

                   _
                   false))
               environment)

    (^ [_ (#.Form (list& [_ (#.Text "lux call")] funcS argsS))])
    (or (contains-self-reference? funcS)
        (list.any? contains-self-reference? argsS))

    (^ [_ (#.Form (list [_ (#.Text "lux let")] register inputS bodyS))])
    (or (contains-self-reference? inputS)
        (contains-self-reference? bodyS))

    (^ [_ (#.Form (list [_ (#.Text "lux if")] inputS thenS elseS))])
    (or (contains-self-reference? inputS)
        (contains-self-reference? thenS)
        (contains-self-reference? elseS))

    (^ [_ (#.Form (list [_ (#.Text "lux loop")] offset [_ (#.Tuple initsS)] bodyS))])
    (or (list.any? contains-self-reference? initsS)
        (contains-self-reference? bodyS))
    
    (^or (^ [_ (#.Form (list& [_ (#.Text "lux recur")] argsS))])
         (^ [_ (#.Form (list& [_ (#.Text procedure)] argsS))]))
    (list.any? contains-self-reference? argsS)

    _
    false
    ))

(def: #export (reify-recursion arity exprS)
  (-> Nat ls.Synthesis ls.Synthesis)
  (loop [exprS exprS]
    (case exprS
      (^ [_ (#.Form (list [_ (#.Text "lux case")] inputS pathS))])
      (` ("lux case" (~ inputS)
          (~ (let [reify-recursion' recur]
               (loop [pathS pathS]
                 (case pathS
                   (^ [_ (#.Form (list [_ (#.Text "lux case alt")] leftS rightS))])
                   (` ("lux case alt" (~ (recur leftS)) (~ (recur rightS))))

                   (^ [_ (#.Form (list [_ (#.Text "lux case seq")] leftS rightS))])
                   (` ("lux case seq" (~ leftS) (~ (recur rightS))))
                   
                   (^ [_ (#.Form (list [_ (#.Text "lux case exec")] bodyS))])
                   (` ("lux case exec" (~ (reify-recursion' bodyS))))
                   
                   _
                   pathS))))))

      (^multi (^ [_ (#.Form (list& [_ (#.Text "lux call")]
                                   [_ (#.Form (list [_ (#.Int 0)]))]
                                   argsS))])
              (n/= arity (list.size argsS)))
      (` ("lux recur" (~+ argsS)))

      (^ [_ (#.Form (list [_ (#.Text "lux let")] register inputS bodyS))])
      (` ("lux let" (~ register) (~ inputS) (~ (recur bodyS))))

      (^ [_ (#.Form (list [_ (#.Text "lux if")] inputS thenS elseS))])
      (` ("lux if" (~ inputS) (~ (recur thenS)) (~ (recur elseS))))

      _
      exprS
      )))

(def: #export (adjust env offset exprS)
  (-> (List Variable) Register ls.Synthesis ls.Synthesis)
  (let [resolve-captured (: (-> Variable Variable)
                            (function (_ var)
                              (let [idx (|> var (i/* -1) int-to-nat n/dec)]
                                (|> env (list.nth idx) maybe.assume))))]
    (loop [exprS exprS]
      (case exprS
        (^code ((~ [_ (#.Nat tag)]) (~ last?) (~ valueS)))
        (` ((~ (code.nat tag)) (~ last?) (~ (recur valueS))))
        
        (^code [(~+ members)])
        (` [(~+ (list/map recur members))])

        (^code ("lux case" (~ inputS) (~ pathS)))
        (` ("lux case" (~ (recur inputS))
            (~ (let [adjust' recur]
                 (loop [pathS pathS]
                   (case pathS
                     (^template [<pattern>]
                       (^ [_ (#.Form (list [_ (#.Text <pattern>)] leftS rightS))])
                       (` (<pattern> (~ (recur leftS)) (~ (recur rightS)))))
                     (["lux case alt"]
                      ["lux case seq"])

                     (^code ("lux case bind" (~ [_ (#.Nat register)])))
                     (` ("lux case bind" (~ (code.nat (n/+ offset register)))))

                     (^ [_ (#.Form (list [_ (#.Text "lux case exec")] bodyS))])
                     (` ("lux case exec" (~ (adjust' bodyS))))
                     
                     _
                     pathS))))))

        (^code ("lux function" (~ arity) [(~+ environment)] (~ bodyS)))
        (` ("lux function" (~ arity)
            [(~+ (list/map (function (_ _var)
                             (case _var
                               (^ [_ (#.Form (list [_ (#.Int var)]))])
                               (` ((~ (code.int (resolve-captured var)))))

                               _
                               _var))
                           environment))]
            (~ bodyS)))

        (^ [_ (#.Form (list& [_ (#.Text "lux call")] funcS argsS))])
        (` ("lux call" (~ (recur funcS)) (~+ (list/map recur argsS))))

        (^ [_ (#.Form (list& [_ (#.Text "lux recur")] argsS))])
        (` ("lux recur" (~+ (list/map recur argsS))))
        
        (^code ("lux let" (~ [_ (#.Nat register)]) (~ inputS) (~ bodyS)))
        (` ("lux let" (~ (code.nat (n/+ offset register)))
            (~ (recur inputS))
            (~ (recur bodyS))))

        (^ [_ (#.Form (list [_ (#.Text "lux if")] inputS thenS elseS))])
        (` ("lux if" (~ (recur inputS))
            (~ (recur thenS))
            (~ (recur elseS))))

        (^ [_ (#.Form (list [_ (#.Text "lux loop")] [_ (#.Nat loop-offset)] [_ (#.Tuple initsS)] bodyS))])
        (` ("lux loop" (~ (code.nat (n/+ offset loop-offset)))
            [(~+ (list/map recur initsS))]
            (~ (recur bodyS))))

        (^ [_ (#.Form (list [_ (#.Int var)]))])
        (if (variableL.captured? var)
          (` ((~ (code.int (resolve-captured var)))))
          (` ((~ (code.int (|> offset nat-to-int (i/+ var)))))))

        (^ [_ (#.Form (list& [_ (#.Text procedure)] argsS))])
        (` ((~ (code.text procedure)) (~+ (list/map recur argsS))))

        _
        exprS
        ))))