path: root/lux-jvm/source/luxc/lang/translation/jvm/function.lux

(.using
 [library
  [lux {"-" Type Label Primitive function}
   [abstract
    ["[0]" monad {"+" do}]
    ["[0]" enum]]
   [control
    ["[0]" pipe]
    ["[0]" function]]
   [data
    ["[0]" product]
    [text
     ["%" format {"+" format}]]
    [collection
     ["[0]" list ("[1]@[0]" functor monoid)]]]
   [math
    [number
     ["n" nat]
     ["i" int]]]
   [target
    [jvm
     ["[0]" type {"+" Type}
      ["[0]" category {"+" Void Value Return Primitive Object Class Array Var Parameter Method}]]]]
   [tool
    [compiler
     [arity {"+" Arity}]
     ["[0]" phase]
     [reference
      [variable {"+" Register}]]
     [language
      [lux
       [analysis {"+" Environment}]
       [synthesis {"+" Synthesis Abstraction Apply}]
       ["[0]" generation]]]
     [meta
      [archive {"+" Archive}
       ["[0]" unit]]
      ["[0]" cache "_"
       ["[1]" artifact]]]]]]]
 [luxc
  [lang
   [host
    ["$" jvm {"+" Label Inst Def Operation Phase Generator}
     ["[0]" def]
     ["_" inst]]]]]
 ["[0]" //
  ["[1][0]" runtime]
  ["[0]" 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@each _.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.each @ (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@each (.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 (is (-> Nat Nat)
                           (|>> ++ (n.+ env_size)))
        store_capturedI (|> (case env_size
                              0 (list)
                              _ (enum.range n.enum 0 (-- env_size)))
                            (list@each (.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@each (.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@each $.new_label (list.repeated num_partials []))
        over_extent (|> (.int function_arity) (i.- (.int apply_arity)))
        casesI (|> (list@composite @labels (list @default))
                   (list.zipped/2 (enum.range n.enum 0 num_partials))
                   (list@each (.function (_ [stage @label])
                                (let [load_partialsI (if (n.> 0 stage)
                                                       (|> (enum.range n.enum 0 (-- stage))
                                                           (list@each (|>> reference.partial_name (load_fieldI class)))
                                                           _.fuse)
                                                       function.identity)]
                                  (cond (i.= over_extent (.int stage))
                                        (|>> (_.label @label)
                                             (_.ALOAD 0)
                                             (pipe.when [(pipe.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@each (|>> 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@each (.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@each (.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 (is Def
                   (if (poly_arg? arity)
                     (|> (n.min arity //runtime.num_apply_variants)
                         (enum.range n.enum 1)
                         (list@each (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 (is 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 unit.ID) (Generator Abstraction))
  (do [! phase.monad]
    [@begin _.make_label
     dependencies (cache.dependencies archive bodyS)
     [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 dependencies
                                  (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.each @ (generate archive) argsS)
     .let [applyI (|> argsI
                      (list.sub //runtime.num_apply_variants)
                      (list@each (.function (_ subI+)
                                   (|>> (_.CHECKCAST //.$Function)
                                        (_.fuse subI+)
                                        (_.INVOKEVIRTUAL //.$Function //runtime.apply_method (//runtime.apply_signature (list.size subI+))))))
                      _.fuse)]]
    (in (|>> functionI
             applyI))))