(.module:
  [lux (#- Type function)
   [abstract
    ["." monad (#+ do)]]
   [control
    [pipe (#+ when> new>)]
    ["." function]]
   [data
    ["." product]
    [text
     ["%" format (#+ format)]]
    [number
     ["n" nat]
     ["i" int]]
    [collection
     ["." list ("#@." functor monoid)]]]
   [target
    [jvm
     ["." type (#+ Type)
      ["." category (#+ Void Value Return Primitive Object Class Array Var Parameter Method)]]]]
   [tool
    [compiler
     [arity (#+ Arity)]
     ["." phase]
     [reference
      [variable (#+ Register)]]
     [language
      [lux
       [analysis (#+ Environment)]
       [synthesis (#+ Synthesis Abstraction Apply)]
       ["." generation]]]
     [meta
      [archive (#+ Archive)]]]]]
  [luxc
   [lang
    [host
     ["$" jvm (#+ Label Inst Def Operation Phase Generator)
      ["." def]
      ["_" inst]]]]]
  ["." //
   ["#." runtime]
   ["." reference]])

(def: arity-field Text "arity")

(def: (poly-arg? arity)
  (-> Arity Bit)
  (n.> 1 arity))

(def: (captured-args env)
  (-> Environment (List (Type Value)))
  (list.repeat (list.size env) //.$Value))

(def: (init-method env arity)
  (-> Environment Arity (Type Method))
  (if (poly-arg? arity)
    (type.method [(list.concat (list (captured-args env)
                                     (list type.int)
                                     (list.repeat (dec arity) //.$Value)))
                  type.void
                  (list)])
    (type.method [(captured-args env) type.void (list)])))

(def: (implementation-method arity)
  (type.method [(list.repeat arity //.$Value) //.$Value (list)]))

(def: get-amount-of-partialsI
  Inst
  (|>> (_.ALOAD 0)
       (_.GETFIELD //.$Function //runtime.partials-field type.int)))

(def: (load-fieldI class field)
  (-> (Type Class) Text Inst)
  (|>> (_.ALOAD 0)
       (_.GETFIELD class field //.$Value)))

(def: (inputsI start amount)
  (-> Register Nat Inst)
  (|> (list.n/range start (n.+ start (dec amount)))
      (list@map _.ALOAD)
      _.fuse))

(def: (applysI start amount)
  (-> Register Nat Inst)
  (let [max-args (n.min amount //runtime.num-apply-variants)
        later-applysI (if (n.> //runtime.num-apply-variants amount)
                        (applysI (n.+ //runtime.num-apply-variants start) (n.- //runtime.num-apply-variants amount))
                        function.identity)]
    (|>> (_.CHECKCAST //.$Function)
         (inputsI start max-args)
         (_.INVOKEVIRTUAL //.$Function //runtime.apply-method (//runtime.apply-signature max-args))
         later-applysI)))

(def: (inc-intI by)
  (-> Nat Inst)
  (|>> (_.int (.int by))
       _.IADD))

(def: (nullsI amount)
  (-> Nat Inst)
  (|> _.NULL
      (list.repeat amount)
      _.fuse))

(def: (instance archive class arity env)
  (-> Archive (Type Class) Arity Environment (Operation Inst))
  (do {@ phase.monad}
    [captureI+ (monad.map @ (reference.variable archive) env)
     #let [argsI (if (poly-arg? arity)
                   (|> (nullsI (dec arity))
                       (list (_.int +0))
                       _.fuse)
                   function.identity)]]
    (wrap (|>> (_.NEW class)
               _.DUP
               (_.fuse captureI+)
               argsI
               (_.INVOKESPECIAL class "<init>" (init-method env arity))))))

(def: (reset-method return)
  (-> (Type Class) (Type Method))
  (type.method [(list) return (list)]))

(def: (with-reset class arity env)
  (-> (Type Class) Arity Environment Def)
  (def.method #$.Public $.noneM "reset" (reset-method class)
              (if (poly-arg? arity)
                (let [env-size (list.size env)
                      captureI (|> (case env-size
                                     0 (list)
                                     _ (list.n/range 0 (dec env-size)))
                                   (list@map (.function (_ source)
                                               (|>> (_.ALOAD 0)
                                                    (_.GETFIELD class (reference.foreign-name source) //.$Value))))
                                   _.fuse)
                      argsI (|> (nullsI (dec arity))
                                (list (_.int +0))
                                _.fuse)]
                  (|>> (_.NEW class)
                       _.DUP
                       captureI
                       argsI
                       (_.INVOKESPECIAL class "<init>" (init-method env arity))
                       _.ARETURN))
                (|>> (_.ALOAD 0)
                     _.ARETURN))))

(def: (with-implementation arity @begin bodyI)
  (-> Nat Label Inst Def)
  (def.method #$.Public $.strictM "impl" (implementation-method arity)
              (|>> (_.label @begin)
                   bodyI
                   _.ARETURN)))

(def: function-init-method
  (type.method [(list type.int) type.void (list)]))

(def: (function-init arity env-size)
  (-> Arity Nat Inst)
  (if (n.= 1 arity)
    (|>> (_.int +0)
         (_.INVOKESPECIAL //.$Function "<init>" function-init-method))
    (|>> (_.ILOAD (inc env-size))
         (_.INVOKESPECIAL //.$Function "<init>" function-init-method))))

(def: (with-init class env arity)
  (-> (Type Class) Environment Arity Def)
  (let [env-size (list.size env)
        offset-partial (: (-> Nat Nat)
                          (|>> inc (n.+ env-size)))
        store-capturedI (|> (case env-size
                              0 (list)
                              _ (list.n/range 0 (dec env-size)))
                            (list@map (.function (_ register)
                                        (|>> (_.ALOAD 0)
                                             (_.ALOAD (inc register))
                                             (_.PUTFIELD class (reference.foreign-name register) //.$Value))))
                            _.fuse)
        store-partialI (if (poly-arg? arity)
                         (|> (list.n/range 0 (n.- 2 arity))
                             (list@map (.function (_ idx)
                                         (let [register (offset-partial idx)]
                                           (|>> (_.ALOAD 0)
                                                (_.ALOAD (inc register))
                                                (_.PUTFIELD class (reference.partial-name idx) //.$Value)))))
                             _.fuse)
                         function.identity)]
    (def.method #$.Public $.noneM "<init>" (init-method env arity)
                (|>> (_.ALOAD 0)
                     (function-init arity env-size)
                     store-capturedI
                     store-partialI
                     _.RETURN))))

(def: (with-apply class env function-arity @begin bodyI apply-arity)
  (-> (Type Class) Environment Arity Label Inst Arity
      Def)
  (let [num-partials (dec function-arity)
        @default ($.new-label [])
        @labels (list@map $.new-label (list.repeat num-partials []))
        over-extent (|> (.int function-arity) (i.- (.int apply-arity)))
        casesI (|> (list@compose @labels (list @default))
                   (list.zip2 (list.n/range 0 num-partials))
                   (list@map (.function (_ [stage @label])
                               (let [load-partialsI (if (n.> 0 stage)
                                                      (|> (list.n/range 0 (dec stage))
                                                          (list@map (|>> reference.partial-name (load-fieldI class)))
                                                          _.fuse)
                                                      function.identity)]
                                 (cond (i.= over-extent (.int stage))
                                       (|>> (_.label @label)
                                            (_.ALOAD 0)
                                            (when> [(new> (n.> 0 stage) [])]
                                                   [(_.INVOKEVIRTUAL class "reset" (reset-method class))])
                                            load-partialsI
                                            (inputsI 1 apply-arity)
                                            (_.INVOKEVIRTUAL class "impl" (implementation-method function-arity))
                                            _.ARETURN)

                                       (i.> over-extent (.int stage))
                                       (let [args-to-completion (|> function-arity (n.- stage))
                                             args-left (|> apply-arity (n.- args-to-completion))]
                                         (|>> (_.label @label)
                                              (_.ALOAD 0)
                                              (_.INVOKEVIRTUAL class "reset" (reset-method class))
                                              load-partialsI
                                              (inputsI 1 args-to-completion)
                                              (_.INVOKEVIRTUAL class "impl" (implementation-method function-arity))
                                              (applysI (inc args-to-completion) args-left)
                                              _.ARETURN))

                                       ## (i.< over-extent (.int stage))
                                       (let [env-size (list.size env)
                                             load-capturedI (|> (case env-size
                                                                  0 (list)
                                                                  _ (list.n/range 0 (dec env-size)))
                                                                (list@map (|>> reference.foreign-name (load-fieldI class)))
                                                                _.fuse)]
                                         (|>> (_.label @label)
                                              (_.NEW class)
                                              _.DUP
                                              load-capturedI
                                              get-amount-of-partialsI
                                              (inc-intI apply-arity)
                                              load-partialsI
                                              (inputsI 1 apply-arity)
                                              (nullsI (|> num-partials (n.- apply-arity) (n.- stage)))
                                              (_.INVOKESPECIAL class "<init>" (init-method env function-arity))
                                              _.ARETURN))
                                       ))))
                   _.fuse)]
    (def.method #$.Public $.noneM //runtime.apply-method (//runtime.apply-signature apply-arity)
                (|>> get-amount-of-partialsI
                     (_.TABLESWITCH +0 (|> num-partials dec .int)
                                    @default @labels)
                     casesI
                     ))))

(def: #export with-environment
  (-> Environment Def)
  (|>> list.enumerate
       (list@map (.function (_ [env-idx env-source])
                   (def.field #$.Private $.finalF (reference.foreign-name env-idx) //.$Value)))
       def.fuse))

(def: (with-partial arity)
  (-> Arity Def)
  (if (poly-arg? arity)
    (|> (list.n/range 0 (n.- 2 arity))
        (list@map (.function (_ idx)
                    (def.field #$.Private $.finalF (reference.partial-name idx) //.$Value)))
        def.fuse)
    function.identity))

(def: #export (with-function archive @begin class env arity bodyI)
  (-> Archive Label Text Environment Arity Inst
      (Operation [Def Inst]))
  (let [classD (type.class class (list))
        applyD (: Def
                  (if (poly-arg? arity)
                    (|> (n.min arity //runtime.num-apply-variants)
                        (list.n/range 1)
                        (list@map (with-apply classD env arity @begin bodyI))
                        (list& (with-implementation arity @begin bodyI))
                        def.fuse)
                    (def.method #$.Public $.strictM //runtime.apply-method (//runtime.apply-signature 1)
                                (|>> (_.label @begin)
                                     bodyI
                                     _.ARETURN))))
        functionD (: Def
                     (|>> (def.int-field #$.Public ($_ $.++F $.staticF $.finalF) arity-field (.int arity))
                          (with-environment env)
                          (with-partial arity)
                          (with-init classD env arity)
                          (with-reset classD arity env)
                          applyD
                          ))]
    (do phase.monad
      [instanceI (instance archive classD arity env)]
      (wrap [functionD instanceI]))))

(def: #export (function generate archive [env arity bodyS])
  (Generator Abstraction)
  (do phase.monad
    [@begin _.make-label
     [function-context bodyI] (generation.with-new-context archive
                                (generation.with-anchor [@begin 1]
                                  (generate archive bodyS)))
     #let [function-class (//.class-name function-context)]
     [functionD instanceI] (with-function archive @begin function-class env arity bodyI)
     _ (generation.save! true ["" (%.nat (product.right function-context))]
                         [function-class
                          (def.class #$.V1_6 #$.Public $.finalC
                                     function-class (list)
                                     //.$Function (list)
                                     functionD)])]
    (wrap instanceI)))

(def: #export (call generate archive [functionS argsS])
  (Generator Apply)
  (do {@ phase.monad}
    [functionI (generate archive functionS)
     argsI (monad.map @ (generate archive) argsS)
     #let [applyI (|> argsI
                      (list.split-all //runtime.num-apply-variants)
                      (list@map (.function (_ chunkI+)
                                  (|>> (_.CHECKCAST //.$Function)
                                       (_.fuse chunkI+)
                                       (_.INVOKEVIRTUAL //.$Function //runtime.apply-method (//runtime.apply-signature (list.size chunkI+))))))
                      _.fuse)]]
    (wrap (|>> functionI
               applyI))))