(.module:
  [lux (#- function)
   [abstract
    ["." monad (#+ do)]]
   [control
    [pipe (#+ when> new>)]
    ["." function]]
   [data
    ["." text
     format]
    [collection
     ["." list ("#/." functor monoid)]]]
   [tool
    [compiler
     [analysis (#+ Arity)]
     [synthesis (#+ Synthesis Abstraction Apply)]
     ["_." reference (#+ Register Variable)]
     ["." phase
      ["." generation]]]]]
  [luxc
   [lang
    [host
     ["$" jvm (#+ Label Inst Def Operation Phase)
      ["." type]
      ["." def]
      ["_" inst]]]]]
  ["." //
   ["." runtime]
   ["." reference]])


(def: arity-field Text "arity")
(def: $Object $.Type (type.class "java.lang.Object" (list)))

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

(def: (reset-method class)
  (-> Text $.Method)
  (type.method (list) (#.Some (type.class class (list))) (list)))

(def: (captured-args env)
  (-> (List Variable) (List $.Type))
  (list.repeat (list.size env) $Object))

(def: (init-method env arity)
  (-> (List Variable) Arity $.Method)
  (if (poly-arg? arity)
    (type.method (list.concat (list (captured-args env)
                                    (list type.int)
                                    (list.repeat (dec arity) $Object)))
                 #.None
                 (list))
    (type.method (captured-args env) #.None (list))))

(def: (implementation-method arity)
  (type.method (list.repeat arity $Object) (#.Some $Object) (list)))

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

(def: (load-fieldI class field)
  (-> Text Text Inst)
  (|>> (_.ALOAD 0)
       (_.GETFIELD class field $Object)))

(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-class)
         (inputsI start max-args)
         (_.INVOKEVIRTUAL //.function-class runtime.apply-method (runtime.apply-signature max-args) #0)
         later-applysI)))

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

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

(def: (with-captured env)
  (-> (List Variable) Def)
  (|> (list.enumerate env)
      (list/map (.function (_ [env-idx env-source])
                  (def.field #$.Private $.finalF (reference.foreign-name env-idx) $Object)))
      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) $Object)))
        def.fuse)
    function.identity))

(def: (instance class arity env)
  (-> Text Arity (List Variable) (Operation Inst))
  (do phase.monad
    [captureI+ (monad.map @ reference.variable 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) #0)))))

(def: (with-reset class arity env)
  (-> Text Arity (List Variable) 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) $Object))))
                                   _.fuse)
                      argsI (|> (nullsI (dec arity))
                                (list (_.int +0))
                                _.fuse)]
                  (|>> (_.NEW class)
                       _.DUP
                       captureI
                       argsI
                       (_.INVOKESPECIAL class "<init>" (init-method env arity) #0)
                       _.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
  $.Method
  (type.method (list type.int) #.None (list)))

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

(def: (with-init class env arity)
  (-> Text (List Variable) 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) $Object))))
                            _.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) $Object)))))
                             _.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)
  (-> Text (List Variable) Arity Label Inst Arity
      Def)
  (let [num-partials (dec function-arity)
        @default ($.new-label [])
        @labels (list/map $.new-label (list.repeat num-partials []))
        arity-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/= arity-over-extent (.int stage))
                                       (|>> (_.label @label)
                                            (_.ALOAD 0)
                                            (when> [(new> (n/> 0 stage) [])]
                                                   [(_.INVOKEVIRTUAL class "reset" (reset-method class) #0)])
                                            load-partialsI
                                            (inputsI 1 apply-arity)
                                            (_.INVOKEVIRTUAL class "impl" (implementation-method function-arity) #0)
                                            _.ARETURN)

                                       (i/> arity-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) #0)
                                              load-partialsI
                                              (inputsI 1 args-to-completion)
                                              (_.INVOKEVIRTUAL class "impl" (implementation-method function-arity) #0)
                                              (applysI (inc args-to-completion) args-left)
                                              _.ARETURN))

                                       ## (i/< arity-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) #0)
                                              _.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
                     (_.INVOKESTATIC //.runtime-class "apply_fail" (type.method (list) #.None (list)) #0)
                     _.NULL
                     _.ARETURN
                     ))))

(def: #export (with-function @begin class env arity bodyI)
  (-> Label Text (List Variable) Arity Inst
      (Operation [Def Inst]))
  (let [env-size (list.size env)
        applyD (: Def
                  (if (poly-arg? arity)
                    (|> (n/min arity runtime.num-apply-variants)
                        (list.n/range 1)
                        (list/map (with-apply class 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-captured env)
                          (with-partial arity)
                          (with-init class env arity)
                          (with-reset class arity env)
                          applyD
                          ))]
    (do phase.monad
      [instanceI (instance class arity env)]
      (wrap [functionD instanceI]))))

(def: #export (function translate [env arity bodyS])
  (-> Phase Abstraction (Operation Inst))
  (do phase.monad
    [@begin _.make-label
     [function-class bodyI] (generation.with-context
                              (generation.with-anchor [@begin 1]
                                (translate bodyS)))
     [functionD instanceI] (with-function @begin function-class env arity bodyI)
     _ (generation.save! true ["" function-class]
                         [function-class
                          (def.class #$.V1_6 #$.Public $.finalC
                                     function-class (list)
                                     ($.simple-class //.function-class) (list)
                                     functionD)])]
    (wrap instanceI)))

(def: (segment size elems)
  (All [a] (-> Nat (List a) (List (List a))))
  (let [[pre post] (list.split size elems)]
    (if (list.empty? post)
      (list pre)
      (list& pre (segment size post)))))

(def: #export (call translate [functionS argsS])
  (-> Phase Apply (Operation Inst))
  (do phase.monad
    [functionI (translate functionS)
     argsI (monad.map @ translate argsS)
     #let [applyI (|> (segment runtime.num-apply-variants argsI)
                      (list/map (.function (_ chunkI+)
                                  (|>> (_.CHECKCAST //.function-class)
                                       (_.fuse chunkI+)
                                       (_.INVOKEVIRTUAL //.function-class runtime.apply-method (runtime.apply-signature (list.size chunkI+)) #0))))
                      _.fuse)]]
    (wrap (|>> functionI
               applyI))))