(.module:
  [library
   [lux (#- Type function)
    [abstract
     ["." monad (#+ do)]
     ["." enum]]
    [control
     [pipe (#+ when> new>)]
     ["." function]]
    [data
     ["." product]
     [text
      ["%" format (#+ format)]]
     [collection
      ["." list ("#@." functor monoid)]]]
    [math
     [number
      ["n" nat]
      ["i" int]]]
    [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 (#+ Context)]]]
      [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 Bit)
  (n.> 1))

(def: (captured_args env)
  (-> (Environment Synthesis) (List (Type Value)))
  (list.repeated (list.size env) //.$Value))

(def: (init_method env arity)
  (-> (Environment Synthesis) Arity (Type Method))
  (if (poly_arg? arity)
    (type.method [(list)
                  (list.together (list (captured_args env)
                                       (list type.int)
                                       (list.repeated (-- arity) //.$Value)))
                  type.void
                  (list)])
    (type.method [(list) (captured_args env) type.void (list)])))

(def: (implementation_method arity)
  (type.method [(list) (list.repeated 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)
  (|> (enum.range n.enum start (n.+ start (-- 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.repeated amount)
      _.fuse))

(def: (instance generate archive class arity env)
  (-> Phase Archive (Type Class) Arity (Environment Synthesis) (Operation Inst))
  (do {@ phase.monad}
    [captureI+ (monad.map @ (generate archive) env)
     .let [argsI (if (poly_arg? arity)
                   (|> (nullsI (-- arity))
                       (list (_.int +0))
                       _.fuse)
                   function.identity)]]
    (in (|>> (_.NEW class)
             _.DUP
             (_.fuse captureI+)
             argsI
             (_.INVOKESPECIAL class "<init>" (init_method env arity))))))

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

(def: (with_reset class arity env)
  (-> (Type Class) Arity (Environment Synthesis) 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)
                                     _ (enum.range n.enum 0 (-- env_size)))
                                   (list@map (.function (_ source)
                                               (|>> (_.ALOAD 0)
                                                    (_.GETFIELD class (reference.foreign_name source) //.$Value))))
                                   _.fuse)
                      argsI (|> (nullsI (-- 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) (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 (++ env_size))
         (_.INVOKESPECIAL //.$Function "<init>" function_init_method))))

(def: (with_init class env arity)
  (-> (Type Class) (Environment Synthesis) Arity Def)
  (let [env_size (list.size env)
        offset_partial (: (-> Nat Nat)
                          (|>> ++ (n.+ env_size)))
        store_capturedI (|> (case env_size
                              0 (list)
                              _ (enum.range n.enum 0 (-- env_size)))
                            (list@map (.function (_ register)
                                        (|>> (_.ALOAD 0)
                                             (_.ALOAD (++ register))
                                             (_.PUTFIELD class (reference.foreign_name register) //.$Value))))
                            _.fuse)
        store_partialI (if (poly_arg? arity)
                         (|> (enum.range n.enum 0 (n.- 2 arity))
                             (list@map (.function (_ idx)
                                         (let [register (offset_partial idx)]
                                           (|>> (_.ALOAD 0)
                                                (_.ALOAD (++ 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 Synthesis) Arity Label Inst Arity
      Def)
  (let [num_partials (-- function_arity)
        @default ($.new_label [])
        @labels (list@map $.new_label (list.repeated num_partials []))
        over_extent (|> (.int function_arity) (i.- (.int apply_arity)))
        casesI (|> (list@compose @labels (list @default))
                   (list.zipped/2 (enum.range n.enum 0 num_partials))
                   (list@map (.function (_ [stage @label])
                               (let [load_partialsI (if (n.> 0 stage)
                                                      (|> (enum.range n.enum 0 (-- 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 (++ args_to_completion) args_left)
                                              _.ARETURN))

                                       ... (i.< over_extent (.int stage))
                                       (let [env_size (list.size env)
                                             load_capturedI (|> (case env_size
                                                                  0 (list)
                                                                  _ (enum.range n.enum 0 (-- 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 -- .int)
                                    @default @labels)
                     casesI
                     ))))

(def: .public with_environment
  (-> (Environment Synthesis) Def)
  (|>> list.enumeration
       (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)
    (|> (enum.range n.enum 0 (n.- 2 arity))
        (list@map (.function (_ idx)
                    (def.field #$.Private $.finalF (reference.partial_name idx) //.$Value)))
        def.fuse)
    function.identity))

(def: .public (with_function generate archive @begin class env arity bodyI)
  (-> Phase Archive Label Text (Environment Synthesis) Arity Inst
      (Operation [Def Inst]))
  (let [classD (type.class class (list))
        applyD (: Def
                  (if (poly_arg? arity)
                    (|> (n.min arity //runtime.num_apply_variants)
                        (enum.range n.enum 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 generate archive classD arity env)]
      (in [functionD instanceI]))))

(def: .public (function' forced_context generate archive [env arity bodyS])
  (-> (Maybe Context) (Generator Abstraction))
  (do {! phase.monad}
    [@begin _.make_label
     [function_context bodyI] (case forced_context
                                (#.Some function_context)
                                (do !
                                  [without_context (generation.with_anchor [@begin 1]
                                                     (generate archive bodyS))]
                                  (in [function_context
                                       without_context]))
                                
                                #.None
                                (generation.with_new_context archive
                                  (generation.with_anchor [@begin 1]
                                    (generate archive bodyS))))
     .let [function_class (//.class_name function_context)]
     [functionD instanceI] (..with_function generate archive @begin function_class env arity bodyI)
     .let [directive [function_class
                      (def.class #$.V1_6 #$.Public $.finalC
                                 function_class (list)
                                 //.$Function (list)
                                 functionD)]]
     _ (generation.execute! directive)
     _ (case forced_context
         #.None
         (generation.save! (product.right function_context) #.None directive)

         (#.Some function_context)
         (in []))]
    (in instanceI)))

(def: .public function
  (Generator Abstraction)
  (..function' #.None))

(def: .public (call generate archive [functionS argsS])
  (Generator Apply)
  (do {@ phase.monad}
    [functionI (generate archive functionS)
     argsI (monad.map @ (generate archive) argsS)
     .let [applyI (|> argsI
                      (list.sub //runtime.num_apply_variants)
                      (list@map (.function (_ subI+)
                                  (|>> (_.CHECKCAST //.$Function)
                                       (_.fuse subI+)
                                       (_.INVOKEVIRTUAL //.$Function //runtime.apply_method (//runtime.apply_signature (list.size subI+))))))
                      _.fuse)]]
    (in (|>> functionI
             applyI))))