(.module:
  [lux #*
   [ffi (#+ import:)]
   [type (#+ :share)]
   [abstract
    ["." monad (#+ do)]]
   [control
    ["." try (#+ Try)]]
   [data
    [identity (#+ Identity)]
    ["." product]
    [text
     ["%" format (#+ format)]]
    [collection
     ["." list ("#@." fold)]
     ["." dictionary (#+ Dictionary)]
     ["." row (#+ Row) ("#@." functor fold)]]]
   [math
    [number
     ["." nat]]]
   [target
    ["/" jvm]]
   [tool
    [compiler
     ["." phase]
     [language
      [lux
       [synthesis (#+ Synthesis)]
       ["." generation]
       ["." directive]
       [phase
        ["." extension
         ["." bundle]
         [directive
          ["./" lux]]]]]]]]]
  [///
   [host
    ["." jvm (#+ Inst)
     ["_" inst]]]])

(import: org/objectweb/asm/Label
  ["#::."
   (new [])])

(def: (literal literal)
  (-> /.Literal Inst)
  (case literal
    (#/.Boolean value) (_.boolean value)
    (#/.Int value) (_.int value)
    (#/.Long value) (_.long value)
    (#/.Double value) (_.double value)
    (#/.Char value) (_.char value)
    (#/.String value) (_.string value)))

(def: (constant instruction)
  (-> /.Constant Inst)
  (case instruction
    (#/.BIPUSH constant) (_.BIPUSH constant)
    
    (#/.SIPUSH constant) (_.SIPUSH constant)

    #/.ICONST_M1 _.ICONST_M1
    #/.ICONST_0 _.ICONST_0
    #/.ICONST_1 _.ICONST_1
    #/.ICONST_2 _.ICONST_2
    #/.ICONST_3 _.ICONST_3
    #/.ICONST_4 _.ICONST_4
    #/.ICONST_5 _.ICONST_5

    #/.LCONST_0 _.LCONST_0
    #/.LCONST_1 _.LCONST_1
    
    #/.FCONST_0 _.FCONST_0
    #/.FCONST_1 _.FCONST_1
    #/.FCONST_2 _.FCONST_2
    
    #/.DCONST_0 _.DCONST_0
    #/.DCONST_1 _.DCONST_1
    
    #/.ACONST_NULL _.NULL

    (#/.LDC literal)
    (..literal literal)
    ))

(def: (int_arithmetic instruction)
  (-> /.Int_Arithmetic Inst)
  (case instruction
    #/.IADD _.IADD
    #/.ISUB _.ISUB
    #/.IMUL _.IMUL
    #/.IDIV _.IDIV
    #/.IREM _.IREM
    #/.INEG _.INEG))

(def: (long_arithmetic instruction)
  (-> /.Long_Arithmetic Inst)
  (case instruction
    #/.LADD _.LADD
    #/.LSUB _.LSUB
    #/.LMUL _.LMUL
    #/.LDIV _.LDIV
    #/.LREM _.LREM
    #/.LNEG _.LNEG))

(def: (float_arithmetic instruction)
  (-> /.Float_Arithmetic Inst)
  (case instruction
    #/.FADD _.FADD
    #/.FSUB _.FSUB
    #/.FMUL _.FMUL
    #/.FDIV _.FDIV
    #/.FREM _.FREM
    #/.FNEG _.FNEG))

(def: (double_arithmetic instruction)
  (-> /.Double_Arithmetic Inst)
  (case instruction
    #/.DADD _.DADD
    #/.DSUB _.DSUB
    #/.DMUL _.DMUL
    #/.DDIV _.DDIV
    #/.DREM _.DREM
    #/.DNEG _.DNEG))

(def: (arithmetic instruction)
  (-> /.Arithmetic Inst)
  (case instruction
    (#/.Int_Arithmetic int_arithmetic)
    (..int_arithmetic int_arithmetic)
    
    (#/.Long_Arithmetic long_arithmetic)
    (..long_arithmetic long_arithmetic)
    
    (#/.Float_Arithmetic float_arithmetic)
    (..float_arithmetic float_arithmetic)
    
    (#/.Double_Arithmetic double_arithmetic)
    (..double_arithmetic double_arithmetic)))

(def: (int_bitwise instruction)
  (-> /.Int_Bitwise Inst)
  (case instruction
    #/.IOR _.IOR
    #/.IXOR _.IXOR
    #/.IAND _.IAND
    #/.ISHL _.ISHL
    #/.ISHR _.ISHR
    #/.IUSHR _.IUSHR))

(def: (long_bitwise instruction)
  (-> /.Long_Bitwise Inst)
  (case instruction
    #/.LOR _.LOR
    #/.LXOR _.LXOR
    #/.LAND _.LAND
    #/.LSHL _.LSHL
    #/.LSHR _.LSHR
    #/.LUSHR _.LUSHR))

(def: (bitwise instruction)
  (-> /.Bitwise Inst)
  (case instruction
    (#/.Int_Bitwise int_bitwise)
    (..int_bitwise int_bitwise)
    
    (#/.Long_Bitwise long_bitwise)
    (..long_bitwise long_bitwise)))

(def: (conversion instruction)
  (-> /.Conversion Inst)
  (case instruction
    #/.I2B _.I2B
    #/.I2S _.I2S
    #/.I2L _.I2L
    #/.I2F _.I2F
    #/.I2D _.I2D
    #/.I2C _.I2C

    #/.L2I _.L2I
    #/.L2F _.L2F
    #/.L2D _.L2D

    #/.F2I _.F2I
    #/.F2L _.F2L
    #/.F2D _.F2D
    
    #/.D2I _.D2I
    #/.D2L _.D2L
    #/.D2F _.D2F))

(def: (array instruction)
  (-> /.Array Inst)
  (case instruction
    #/.ARRAYLENGTH _.ARRAYLENGTH

    (#/.NEWARRAY type) (_.NEWARRAY type)
    (#/.ANEWARRAY type) (_.ANEWARRAY type)

    #/.BALOAD _.BALOAD
    #/.BASTORE _.BASTORE

    #/.SALOAD _.SALOAD
    #/.SASTORE _.SASTORE

    #/.IALOAD _.IALOAD
    #/.IASTORE _.IASTORE

    #/.LALOAD _.LALOAD
    #/.LASTORE _.LASTORE

    #/.FALOAD _.FALOAD
    #/.FASTORE _.FASTORE

    #/.DALOAD _.DALOAD
    #/.DASTORE _.DASTORE
    
    #/.CALOAD _.CALOAD
    #/.CASTORE _.CASTORE

    #/.AALOAD _.AALOAD
    #/.AASTORE _.AASTORE))

(def: (object instruction)
  (-> /.Object Inst)
  (case instruction
    (^template [<tag> <inst>]
      [(<tag> class field_name field_type)
       (<inst> class field_name field_type)])
    ([#/.GETSTATIC _.GETSTATIC]
     [#/.PUTSTATIC _.PUTSTATIC]
     [#/.GETFIELD _.GETFIELD]
     [#/.PUTFIELD _.PUTFIELD])
    
    (#/.NEW type) (_.NEW type)
    
    (#/.INSTANCEOF type) (_.INSTANCEOF type)
    (#/.CHECKCAST type) (_.CHECKCAST type)

    (^template [<tag> <inst>]
      [(<tag> class method_name method_type)
       (<inst> class method_name method_type)])
    ([#/.INVOKEINTERFACE _.INVOKEINTERFACE]
     [#/.INVOKESPECIAL _.INVOKESPECIAL]
     [#/.INVOKESTATIC _.INVOKESTATIC]
     [#/.INVOKEVIRTUAL _.INVOKEVIRTUAL])
    ))

(def: (local_int instruction)
  (-> /.Local_Int Inst)
  (case instruction
    (#/.ILOAD register) (_.ILOAD register)
    (#/.ISTORE register) (_.ISTORE register)))

(def: (local_long instruction)
  (-> /.Local_Long Inst)
  (case instruction
    (#/.LLOAD register) (_.LLOAD register)
    (#/.LSTORE register) (_.LSTORE register)))

(def: (local_float instruction)
  (-> /.Local_Float Inst)
  (case instruction
    (#/.FLOAD register) (_.FLOAD register)
    (#/.FSTORE register) (_.FSTORE register)))

(def: (local_double instruction)
  (-> /.Local_Double Inst)
  (case instruction
    (#/.DLOAD register) (_.DLOAD register)
    (#/.DSTORE register) (_.DSTORE register)))

(def: (local_object instruction)
  (-> /.Local_Object Inst)
  (case instruction
    (#/.ALOAD register) (_.ALOAD register)
    (#/.ASTORE register) (_.ASTORE register)))

(def: (local instruction)
  (-> /.Local Inst)
  (case instruction
    (#/.Local_Int instruction) (..local_int instruction)
    (#/.IINC register) (_.IINC register)
    (#/.Local_Long instruction) (..local_long instruction)
    (#/.Local_Float instruction) (..local_float instruction)
    (#/.Local_Double instruction) (..local_double instruction)
    (#/.Local_Object instruction) (..local_object instruction)))

(def: (stack instruction)
  (-> /.Stack Inst)
  (case instruction
    #/.DUP _.DUP
    #/.DUP_X1 _.DUP_X1
    #/.DUP_X2 _.DUP_X2
    #/.DUP2 _.DUP2
    #/.DUP2_X1 _.DUP2_X1
    #/.DUP2_X2 _.DUP2_X2
    #/.SWAP _.SWAP
    #/.POP _.POP
    #/.POP2 _.POP2))

(def: (comparison instruction)
  (-> /.Comparison Inst)
  (case instruction
    #/.LCMP _.LCMP
    
    #/.FCMPG _.FCMPG
    #/.FCMPL _.FCMPL

    #/.DCMPG _.DCMPG
    #/.DCMPL _.DCMPL))

(def: (branching instruction)
  (-> (/.Branching org/objectweb/asm/Label) Inst)
  (case instruction
    (#/.IF_ICMPEQ label) (_.IF_ICMPEQ label)
    (#/.IF_ICMPGE label) (_.IF_ICMPGE label)
    (#/.IF_ICMPGT label) (_.IF_ICMPGT label)
    (#/.IF_ICMPLE label) (_.IF_ICMPLE label)
    (#/.IF_ICMPLT label) (_.IF_ICMPLT label)
    (#/.IF_ICMPNE label) (_.IF_ICMPNE label)
    (#/.IFEQ label) (_.IFEQ label)
    (#/.IFGE label) (_.IFGE label)
    (#/.IFGT label) (_.IFGT label)
    (#/.IFLE label) (_.IFLE label)
    (#/.IFLT label) (_.IFLT label)
    (#/.IFNE label) (_.IFNE label)

    (#/.TABLESWITCH min max default labels)
    (_.TABLESWITCH min max default labels)
    
    (#/.LOOKUPSWITCH default keys+labels)
    (_.LOOKUPSWITCH default keys+labels)

    (#/.IF_ACMPEQ label) (_.IF_ACMPEQ label)
    (#/.IF_ACMPNE label) (_.IF_ACMPNE label)
    (#/.IFNONNULL label) (_.IFNONNULL label)
    (#/.IFNULL label) (_.IFNULL label)))

(def: (exception instruction)
  (-> (/.Exception org/objectweb/asm/Label) Inst)
  (case instruction
    (#/.Try start end handler exception) (_.try start end handler exception)
    #/.ATHROW _.ATHROW))

(def: (concurrency instruction)
  (-> /.Concurrency Inst)
  (case instruction
    #/.MONITORENTER _.MONITORENTER
    #/.MONITOREXIT _.MONITOREXIT))

(def: (return instruction)
  (-> /.Return Inst)
  (case instruction
    #/.RETURN _.RETURN
    #/.IRETURN _.IRETURN
    #/.LRETURN _.LRETURN
    #/.FRETURN _.FRETURN
    #/.DRETURN _.DRETURN
    #/.ARETURN _.ARETURN))

(def: (control instruction)
  (-> (/.Control org/objectweb/asm/Label) Inst)
  (case instruction
    (#/.GOTO label) (_.GOTO label)
    (#/.Branching instruction) (..branching instruction)
    (#/.Exception instruction) (..exception instruction)
    (#/.Concurrency instruction) (..concurrency instruction)
    (#/.Return instruction) (..return instruction)))

(def: (instruction instruction)
  (-> (/.Instruction Inst org/objectweb/asm/Label) Inst)
  (case instruction
    #/.NOP _.NOP
    (#/.Constant instruction) (..constant instruction)
    (#/.Arithmetic instruction) (..arithmetic instruction)
    (#/.Bitwise instruction) (..bitwise instruction)
    (#/.Conversion instruction) (..conversion instruction)
    (#/.Array instruction) (..array instruction)
    (#/.Object instruction) (..object instruction)
    (#/.Local instruction) (..local instruction)
    (#/.Stack instruction) (..stack instruction)
    (#/.Comparison instruction) (..comparison instruction)
    (#/.Control instruction) (..control instruction)
    (#/.Embedded embedded) embedded))

(type: Mapping
  (Dictionary /.Label org/objectweb/asm/Label))

(type: (Re_labeler context)
  (-> [Mapping (context /.Label)]
      [Mapping (context org/objectweb/asm/Label)]))

(def: (relabel [mapping label])
  (Re_labeler Identity)
  (case (dictionary.get label mapping)
    (#.Some label)
    [mapping label]

    #.None
    (let [label' (org/objectweb/asm/Label::new)]
      [(dictionary.put label label' mapping) label'])))

(def: (relabel_branching [mapping instruction])
  (Re_labeler /.Branching)
  (case instruction
    (^template [<tag>]
      [(<tag> label)
       (let [[mapping label] (..relabel [mapping label])]
         [mapping (<tag> label)])])
    ([#/.IF_ICMPEQ] [#/.IF_ICMPGE] [#/.IF_ICMPGT] [#/.IF_ICMPLE] [#/.IF_ICMPLT] [#/.IF_ICMPNE]
     [#/.IFEQ] [#/.IFNE] [#/.IFGE] [#/.IFGT] [#/.IFLE] [#/.IFLT]

     [#/.IF_ACMPEQ] [#/.IF_ACMPNE] [#/.IFNONNULL] [#/.IFNULL])

    (#/.TABLESWITCH min max default labels)
    (let [[mapping default] (..relabel [mapping default])
          [mapping labels] (list@fold (function (_ input [mapping output])
                                        (let [[mapping input] (..relabel [mapping input])]
                                          [mapping (list& input output)]))
                                      [mapping (list)] labels)]
      [mapping (#/.TABLESWITCH min max default (list.reverse labels))])
    
    (#/.LOOKUPSWITCH default keys+labels)
    (let [[mapping default] (..relabel [mapping default])
          [mapping keys+labels] (list@fold (function (_ [expected input] [mapping output])
                                             (let [[mapping input] (..relabel [mapping input])]
                                               [mapping (list& [expected input] output)]))
                                           [mapping (list)] keys+labels)]
      [mapping (#/.LOOKUPSWITCH default (list.reverse keys+labels))])
    ))

(def: (relabel_exception [mapping instruction])
  (Re_labeler /.Exception)
  (case instruction
    (#/.Try start end handler exception)
    (let [[mapping start] (..relabel [mapping start])
          [mapping end] (..relabel [mapping end])
          [mapping handler] (..relabel [mapping handler])]
      [mapping (#/.Try start end handler exception)])
    
    #/.ATHROW
    [mapping #/.ATHROW]
    ))

(def: (relabel_control [mapping instruction])
  (Re_labeler /.Control)
  (case instruction
    (^template [<tag> <relabel>]
      [(<tag> instruction)
       (let [[mapping instruction] (<relabel> [mapping instruction])]
         [mapping (<tag> instruction)])])
    ([#/.GOTO ..relabel]
     [#/.Branching ..relabel_branching]
     [#/.Exception ..relabel_exception])

    (^template [<tag>]
      [(<tag> instruction)
       [mapping (<tag> instruction)]])
    ([#/.Concurrency] [#/.Return])
    ))

(def: (relabel_instruction [mapping instruction])
  (Re_labeler (/.Instruction Inst))
  (case instruction
    (#/.Embedded embedded)
    [mapping (#/.Embedded embedded)]

    #/.NOP
    [mapping #/.NOP]

    (^template [<tag>]
      [(<tag> instruction)
       [mapping (<tag> instruction)]])
    ([#/.Constant]
     [#/.Arithmetic]
     [#/.Bitwise]
     [#/.Conversion]
     [#/.Array]
     [#/.Object]
     [#/.Local]
     [#/.Stack]
     [#/.Comparison])
    
    (#/.Control instruction)
    (let [[mapping instruction] (..relabel_control [mapping instruction])]
      [mapping (#/.Control instruction)])))

(def: (relabel_bytecode [mapping bytecode])
  (Re_labeler (/.Bytecode Inst))
  (row@fold (function (_ input [mapping output])
              (let [[mapping input'] (..relabel_instruction [mapping input])]
                [mapping (row.add input' output)]))
            [mapping (row.row)]
            bytecode))

(def: fresh
  Mapping
  (dictionary.new nat.hash))

(def: bytecode
  (-> (/.Bytecode Inst /.Label) jvm.Inst)
  (|>> [..fresh]
       ..relabel_bytecode
       product.right
       (row@map ..instruction)
       row.to_list
       _.fuse))

(with_expansions [<anchor> (as_is jvm.Anchor)
                  <expression> (as_is Inst)
                  <directive> (as_is jvm.Definition)
                  <type_vars> (as_is <anchor> <expression> <directive>)]
  (type: Handler
    ## (generation.Handler jvm.Anchor (/.Bytecode Inst /.Label) jvm.Definition)
    (-> extension.Name
        (phase.Phase [(extension.Bundle <type_vars>)
                      (generation.State <type_vars>)]
                     Synthesis
                     <expression>)
        (phase.Phase [(extension.Bundle <type_vars>)
                      (generation.State <type_vars>)]
                     (List Synthesis)
                     (/.Bytecode Inst /.Label)))))

(def: (true_handler extender pseudo)
  (-> jvm.Extender Any jvm.Handler)
  (function (_ extension_name phase archive inputs)
    (\ phase.monad map
       (|>> (:as (/.Bytecode Inst /.Label)) ..bytecode)
       ((extender pseudo) extension_name phase archive inputs))))

(def: (def::generation extender)
  (-> jvm.Extender
      (directive.Handler jvm.Anchor jvm.Inst jvm.Definition))
  (function (handler extension_name phase archive inputsC+)
    (case inputsC+
      (^ (list nameC valueC))
      (do phase.monad
        [[_ _ name] (lux/.evaluate! archive Text nameC)
         [_ handlerV] (lux/.generator archive (:as Text name) ..Handler valueC)
         _ (|> handlerV
               (..true_handler extender)
               (extension.install extender (:as Text name))
               directive.lift_generation)
         _ (directive.lift_generation
            (generation.log! (format "Generation " (%.text (:as Text name)))))]
        (wrap directive.no_requirements))

      _
      (phase.throw extension.invalid_syntax [extension_name %.code inputsC+]))))

(def: #export (bundle extender)
  (-> jvm.Extender
      (directive.Bundle jvm.Anchor jvm.Inst jvm.Definition))
  (|> bundle.empty
      (dictionary.put "lux def generation" (..def::generation extender))))