(.module: {#.doc "Type-checking functionality."}
  [lux #*
   [abstract
    [functor (#+ Functor)]
    [apply (#+ Apply)]
    ["." monad (#+ Monad do)]]
   [control
    ["." try (#+ Try)]
    ["." exception (#+ Exception exception:)]]
   [data
    ["." maybe]
    ["." product]
    [number
     ["n" nat ("#\." decimal)]]
    ["." text ("#\." monoid equivalence)]
    [collection
     ["." list]
     ["." set (#+ Set)]]]]
  ["." // ("#\." equivalence)])

(template: (!n/= reference subject)
  ("lux i64 =" reference subject))

(template: (!text\= reference subject)
  ("lux text =" reference subject))

(exception: #export (unknown_type_var {id Nat})
  (exception.report
   ["ID" (n\encode id)]))

(exception: #export (unbound_type_var {id Nat})
  (exception.report
   ["ID" (n\encode id)]))

(exception: #export (invalid_type_application {funcT Type} {argT Type})
  (exception.report
   ["Type function" (//.format funcT)]
   ["Type argument" (//.format argT)]))

(exception: #export (cannot_rebind_var {id Nat} {type Type} {bound Type})
  (exception.report
   ["Var" (n\encode id)]
   ["Wanted Type" (//.format type)]
   ["Current Type" (//.format bound)]))

(exception: #export (type_check_failed {expected Type} {actual Type})
  (exception.report
   ["Expected" (//.format expected)]
   ["Actual" (//.format actual)]))

(type: #export Var Nat)

(type: #export Assumption [Type Type])

(type: #export (Check a)
  (-> Type_Context (Try [Type_Context a])))

(type: #export (Checker a)
  (-> (List Assumption) a a (Check (List Assumption))))

(type: #export Type_Vars
  (List [Var (Maybe Type)]))

(structure: #export functor
  (Functor Check)
  
  (def: (map f fa)
    (function (_ context)
      (case (fa context)
        (#try.Success [context' output])
        (#try.Success [context' (f output)])

        (#try.Failure error)
        (#try.Failure error)))))

(structure: #export apply
  (Apply Check)
  
  (def: &functor ..functor)

  (def: (apply ff fa)
    (function (_ context)
      (case (ff context)
        (#try.Success [context' f])
        (case (fa context')
          (#try.Success [context'' a])
          (#try.Success [context'' (f a)])

          (#try.Failure error)
          (#try.Failure error))

        (#try.Failure error)
        (#try.Failure error)
        )))
  )

(structure: #export monad
  (Monad Check)
  
  (def: &functor ..functor)

  (def: (wrap x)
    (function (_ context)
      (#try.Success [context x])))

  (def: (join ffa)
    (function (_ context)
      (case (ffa context)
        (#try.Success [context' fa])
        (case (fa context')
          (#try.Success [context'' a])
          (#try.Success [context'' a])

          (#try.Failure error)
          (#try.Failure error))

        (#try.Failure error)
        (#try.Failure error)
        )))
  )

(open: "check\." ..monad)

(def: (var::new id plist)
  (-> Var Type_Vars Type_Vars)
  (#.Cons [id #.None] plist))

(def: (var::get id plist)
  (-> Var Type_Vars (Maybe (Maybe Type)))
  (case plist
    (#.Cons [var_id var_type]
            plist')
    (if (!n/= id var_id)
      (#.Some var_type)
      (var::get id plist'))

    #.Nil
    #.None))

(def: (var::put id value plist)
  (-> Var (Maybe Type) Type_Vars Type_Vars)
  (case plist
    #.Nil
    (list [id value])

    (#.Cons [var_id var_type]
            plist')
    (if (!n/= id var_id)
      (#.Cons [var_id value]
              plist')
      (#.Cons [var_id var_type]
              (var::put id value plist')))))

(def: #export (run context proc)
  (All [a] (-> Type_Context (Check a) (Try a)))
  (case (proc context)
    (#try.Success [context' output])
    (#try.Success output)

    (#try.Failure error)
    (#try.Failure error)))

(def: #export (fail message)
  (All [a] (-> Text (Check a)))
  (function (_ context)
    (#try.Failure message)))

(def: #export (assert message test)
  (-> Text Bit (Check Any))
  (function (_ context)
    (if test
      (#try.Success [context []])
      (#try.Failure message))))

(def: #export (throw exception message)
  (All [e a] (-> (Exception e) e (Check a)))
  (..fail (exception.construct exception message)))

(def: #export existential
  {#.doc "A producer of existential types."}
  (Check [Nat Type])
  (function (_ context)
    (let [id (get@ #.ex_counter context)]
      (#try.Success [(update@ #.ex_counter inc context)
                     [id (#.Ex id)]]))))

(template [<name> <outputT> <fail> <succeed>]
  [(def: #export (<name> id)
     (-> Var (Check <outputT>))
     (function (_ context)
       (case (|> context (get@ #.var_bindings) (var::get id))
         (^or (#.Some (#.Some (#.Var _)))
              (#.Some #.None))
         (#try.Success [context <fail>])
         
         (#.Some (#.Some bound))
         (#try.Success [context <succeed>])

         #.None
         (exception.throw ..unknown_type_var id))))]

  [bound? Bit          false  true]
  [read   (Maybe Type) #.None (#.Some bound)]
  )

(def: #export (read! id)
  (-> Var (Check Type))
  (do ..monad
    [?type (read id)]
    (case ?type
      (#.Some type)
      (wrap type)

      #.None
      (..throw ..unbound_type_var id))))

(def: (peek id)
  (-> Var (Check Type))
  (function (_ context)
    (case (|> context (get@ #.var_bindings) (var::get id))
      (#.Some (#.Some bound))
      (#try.Success [context bound])

      (#.Some _)
      (exception.throw ..unbound_type_var id)

      _
      (exception.throw ..unknown_type_var id))))

(def: #export (bind type id)
  (-> Type Var (Check Any))
  (function (_ context)
    (case (|> context (get@ #.var_bindings) (var::get id))
      (#.Some #.None)
      (#try.Success [(update@ #.var_bindings (var::put id (#.Some type)) context)
                     []])

      (#.Some (#.Some bound))
      (exception.throw ..cannot_rebind_var [id type bound])
      
      _
      (exception.throw ..unknown_type_var id))))

(def: (update type id)
  (-> Type Var (Check Any))
  (function (_ context)
    (case (|> context (get@ #.var_bindings) (var::get id))
      (#.Some _)
      (#try.Success [(update@ #.var_bindings (var::put id (#.Some type)) context)
                     []])
      
      _
      (exception.throw ..unknown_type_var id))))

(def: #export var
  (Check [Var Type])
  (function (_ context)
    (let [id (get@ #.var_counter context)]
      (#try.Success [(|> context
                         (update@ #.var_counter inc)
                         (update@ #.var_bindings (var::new id)))
                     [id (#.Var id)]]))))

(def: (apply_type! funcT argT)
  (-> Type Type (Check Type))
  (case funcT
    (#.Var func_id)
    (do ..monad
      [?funcT' (read func_id)]
      (case ?funcT'
        (#.Some funcT')
        (apply_type! funcT' argT)

        _
        (throw ..invalid_type_application [funcT argT])))

    (#.Apply argT' funcT')
    (do ..monad
      [funcT'' (apply_type! funcT' argT')]
      (apply_type! funcT'' argT))

    _
    (case (//.apply (list argT) funcT)
      (#.Some output)
      (check\wrap output)

      _
      (throw ..invalid_type_application [funcT argT]))))

(type: #export Ring (Set Var))

(def: empty_ring Ring (set.new n.hash))

## TODO: Optimize this by not using sets anymore.
(def: #export (ring start)
  (-> Var (Check Ring))
  (function (_ context)
    (loop [current start
           output (set.add start empty_ring)]
      (case (|> context (get@ #.var_bindings) (var::get current))
        (#.Some (#.Some type))
        (case type
          (#.Var post)
          (if (!n/= start post)
            (#try.Success [context output])
            (recur post (set.add post output)))
          
          _
          (#try.Success [context empty_ring]))

        (#.Some #.None)
        (#try.Success [context output])
        
        #.None
        (exception.throw ..unknown_type_var current)))))

(def: #export fresh_context
  Type_Context
  {#.var_counter 0
   #.ex_counter 0
   #.var_bindings (list)})

(def: (attempt op)
  (All [a] (-> (Check a) (Check (Maybe a))))
  (function (_ context)
    (case (op context)
      (#try.Success [context' output])
      (#try.Success [context' (#.Some output)])

      (#try.Failure _)
      (#try.Success [context #.None]))))

(def: (either left right)
  (All [a] (-> (Check a) (Check a) (Check a)))
  (function (_ context)
    (case (left context)
      (#try.Failure _)
      (right context)

      output
      output)))

(def: (assumed? [e a] assumptions)
  (-> Assumption (List Assumption) Bit)
  (list.any? (function (_ [e' a'])
               (and (//\= e e')
                    (//\= a a')))
             assumptions))

(def: (assume! assumption assumptions)
  (-> Assumption (List Assumption) (List Assumption))
  (#.Cons assumption assumptions))

## TODO: "if_bind" can be optimized...
(def: (if_bind id type then else)
  (All [a]
    (-> Var Type (Check a) (-> Type (Check a))
        (Check a)))
  ($_ either
      (do ..monad
        [_ (..bind type id)]
        then)
      (do {! ..monad}
        [ring (..ring id)
         _ (assert "" (n.> 1 (set.size ring)))
         _ (monad.map ! (update type) (set.to_list ring))]
        then)
      (do ..monad
        [?bound (read id)]
        (else (maybe.default (#.Var id) ?bound)))))

## TODO: "link_2" can be optimized...
(def: (link_2 left right)
  (-> Var Var (Check Any))
  (do ..monad
    [_ (..bind (#.Var right) left)]
    (..bind (#.Var left) right)))

## TODO: "link_3" can be optimized...
(def: (link_3 interpose to from)
  (-> Var Var Var (Check Any))
  (do ..monad
    [_ (update (#.Var interpose) from)]
    (update (#.Var to) interpose)))

## TODO: "check_vars" can be optimized...
(def: (check_vars check' assumptions idE idA)
  (-> (Checker Type) (Checker Var))
  (if (!n/= idE idA)
    (check\wrap assumptions)
    (do {! ..monad}
      [ebound (attempt (peek idE))
       abound (attempt (peek idA))]
      (case [ebound abound]
        ## Link the 2 variables circularly
        [#.None #.None]
        (do !
          [_ (link_2 idE idA)]
          (wrap assumptions))

        ## Interpose new variable between 2 existing links
        [(#.Some etype) #.None]
        (case etype
          (#.Var targetE)
          (do !
            [_ (link_3 idA targetE idE)]
            (wrap assumptions))

          _
          (check' assumptions etype (#.Var idA)))

        ## Interpose new variable between 2 existing links
        [#.None (#.Some atype)]
        (case atype
          (#.Var targetA)
          (do !
            [_ (link_3 idE targetA idA)]
            (wrap assumptions))

          _
          (check' assumptions (#.Var idE) atype))

        [(#.Some etype) (#.Some atype)]
        (case [etype atype]
          [(#.Var targetE) (#.Var targetA)]
          (do !
            [ringE (..ring idE)
             ringA (..ring idA)]
            (if (\ set.equivalence = ringE ringA)
              (wrap assumptions)
              ## Fuse 2 rings
              (do !
                [_ (monad.fold ! (function (_ interpose to)
                                   (do !
                                     [_ (link_3 interpose to idE)]
                                     (wrap interpose)))
                               targetE
                               (set.to_list ringA))]
                (wrap assumptions))))

          (^template [<pattern> <id> <type>]
            [<pattern>
             (do !
               [ring (..ring <id>)
                _ (monad.map ! (update <type>) (set.to_list ring))]
               (wrap assumptions))])
          ([[(#.Var _) _] idE atype]
           [[_ (#.Var _)] idA etype])
          
          _
          (check' assumptions etype atype))))))

## TODO: "check_apply" can be optimized...
(def: (check_apply check' assumptions expected actual)
  (-> (Checker Type) (Checker [Type Type]))
  (let [[expected_input expected_function] expected
        [actual_input actual_function] actual]
    (case [expected_function actual_function]
      [(#.Ex exE) (#.Ex exA)]
      (if (!n/= exE exA)
        (check' assumptions expected_input actual_input)
        (fail ""))

      [(#.UnivQ _ _) (#.Ex _)]
      (do ..monad
        [expected' (apply_type! expected_function expected_input)]
        (check' assumptions expected' (#.Apply actual)))

      [(#.Ex _) (#.UnivQ _ _)]
      (do ..monad
        [actual' (apply_type! actual_function actual_input)]
        (check' assumptions (#.Apply expected) actual'))

      [(#.Apply [expected_input' expected_function']) (#.Ex _)]
      (do ..monad
        [expected_function'' (apply_type! expected_function' expected_input')]
        (check' assumptions (#.Apply [expected_input expected_function'']) (#.Apply actual)))

      [(#.Ex _) (#.Apply [actual_input' actual_function'])]
      (do ..monad
        [actual_function'' (apply_type! actual_function' actual_input')]
        (check' assumptions (#.Apply expected) (#.Apply [actual_input actual_function''])))

      (^or [(#.Ex _) _] [_ (#.Ex _)])
      (do ..monad
        [assumptions (check' assumptions expected_function actual_function)]
        (check' assumptions expected_input actual_input))

      [(#.Var id) _]
      (function (_ context)
        (case ((do ..monad
                 [expected_function' (read! id)]
                 (check' assumptions (#.Apply expected_input expected_function') (#.Apply actual)))
               context)
          (#try.Success output)
          (#try.Success output)

          (#try.Failure _)
          (case actual_function
            (#.UnivQ _ _)
            ((do ..monad
               [actual' (apply_type! actual_function actual_input)]
               (check' assumptions (#.Apply expected) actual'))
             context)
            
            (#.Ex exA)
            ((do ..monad
               [assumptions (check' assumptions expected_function actual_function)]
               (check' assumptions expected_input actual_input))
             context)

            _
            ((do ..monad
               [assumptions (check' assumptions expected_function actual_function)
                expected' (apply_type! actual_function expected_input)
                actual' (apply_type! actual_function actual_input)]
               (check' assumptions expected' actual'))
             context))))

      [_ (#.Var id)]
      (function (_ context)
        (case ((do ..monad
                 [actual_function' (read! id)]
                 (check' assumptions (#.Apply expected) (#.Apply actual_input actual_function')))
               context)
          (#try.Success output)
          (#try.Success output)

          _
          ((do ..monad
             [assumptions (check' assumptions expected_function actual_function)
              expected' (apply_type! expected_function expected_input)
              actual' (apply_type! expected_function actual_input)]
             (check' assumptions expected' actual'))
           context)))

      _
      (fail ""))))

(def: (with exception parameter check)
  (All [e a] (-> (Exception e) e (Check a) (Check a)))
  (|>> check (exception.with exception parameter)))

## TODO: "check'" can be optimized...
(def: #export (check' assumptions expected actual)
  {#.doc "Type-check to ensure that the 'expected' type subsumes the 'actual' type."}
  (Checker Type)
  (if (is? expected actual)
    (check\wrap assumptions)
    (with type_check_failed [expected actual]
      (case [expected actual]
        [(#.Var idE) (#.Var idA)]
        (check_vars check' assumptions idE idA)
        
        [(#.Var id) _]
        (if_bind id actual
                 (check\wrap assumptions)
                 (function (_ bound)
                   (check' assumptions bound actual)))
        
        [_ (#.Var id)]
        (if_bind id expected
                 (check\wrap assumptions)
                 (function (_ bound)
                   (check' assumptions expected bound)))

        (^template [<fE> <fA>]
          [[(#.Apply aE <fE>) (#.Apply aA <fA>)]
           (check_apply check' assumptions [aE <fE>] [aA <fA>])])
        ([F1 (#.Ex ex)]
         [(#.Ex exE) fA]
         [fE (#.Var idA)]
         [(#.Var idE) fA])
        
        [(#.Apply A F) _]
        (let [new_assumption [expected actual]]
          (if (assumed? new_assumption assumptions)
            (check\wrap assumptions)
            (do ..monad
              [expected' (apply_type! F A)]
              (check' (assume! new_assumption assumptions) expected' actual))))

        [_ (#.Apply A F)]
        (do ..monad
          [actual' (apply_type! F A)]
          (check' assumptions expected actual'))

        ## TODO: Refactor-away as cold-code
        (^template [<tag> <instancer>]
          [[(<tag> _) _]
           (do ..monad
             [[_ paramT] <instancer>
              expected' (apply_type! expected paramT)]
             (check' assumptions expected' actual))])
        ([#.UnivQ ..existential]
         [#.ExQ ..var])

        ## TODO: Refactor-away as cold-code
        (^template [<tag> <instancer>]
          [[_ (<tag> _)]
           (do ..monad
             [[_ paramT] <instancer>
              actual' (apply_type! actual paramT)]
             (check' assumptions expected actual'))])
        ([#.UnivQ ..var]
         [#.ExQ ..existential])

        [(#.Primitive e_name e_params) (#.Primitive a_name a_params)]
        (if (!text\= e_name a_name)
          (loop [assumptions assumptions
                 e_params e_params
                 a_params a_params]
            (case [e_params a_params]
              [#.Nil #.Nil]
              (check\wrap assumptions)
              
              [(#.Cons e_head e_tail) (#.Cons a_head a_tail)]
              (do ..monad
                [assumptions' (check' assumptions e_head a_head)]
                (recur assumptions' e_tail a_tail))

              _
              (fail "")))
          (fail ""))

        (^template [<compose>]
          [[(<compose> eL eR) (<compose> aL aR)]
           (do ..monad
             [assumptions (check' assumptions eL aL)]
             (check' assumptions eR aR))])
        ([#.Sum]
         [#.Product])
        
        [(#.Function eI eO) (#.Function aI aO)]
        (do ..monad
          [assumptions (check' assumptions aI eI)]
          (check' assumptions eO aO))

        [(#.Ex e!id) (#.Ex a!id)]
        (if (!n/= e!id a!id)
          (check\wrap assumptions)
          (fail ""))

        [(#.Named _ ?etype) _]
        (check' assumptions ?etype actual)

        [_ (#.Named _ ?atype)]
        (check' assumptions expected ?atype)

        _
        (fail "")))))

(def: #export (check expected actual)
  {#.doc "Type-check to ensure that the 'expected' type subsumes the 'actual' type."}
  (-> Type Type (Check Any))
  (check' (list) expected actual))

(def: #export (checks? expected actual)
  {#.doc "A simple type-checking function that just returns a yes/no answer."}
  (-> Type Type Bit)
  (case (run fresh_context (check' (list) expected actual))
    (#try.Failure _)
    false

    (#try.Success _)
    true))

(def: #export context
  (Check Type_Context)
  (function (_ context)
    (#try.Success [context context])))

(def: #export (clean inputT)
  (-> Type (Check Type))
  (case inputT
    (#.Primitive name paramsT+)
    (|> paramsT+
        (monad.map ..monad clean)
        (check\map (|>> (#.Primitive name))))

    (^or (#.Parameter _) (#.Ex _) (#.Named _))
    (check\wrap inputT)

    (^template [<tag>]
      [(<tag> leftT rightT)
       (do ..monad
         [leftT' (clean leftT)]
         (|> (clean rightT)
             (check\map (|>> (<tag> leftT')))))])
    ([#.Sum] [#.Product] [#.Function] [#.Apply])

    (#.Var id)
    (do ..monad
      [?actualT (read id)]
      (case ?actualT
        (#.Some actualT)
        (clean actualT)

        _
        (wrap inputT)))

    (^template [<tag>]
      [(<tag> envT+ unquantifiedT)
       (do {! ..monad}
         [envT+' (monad.map ! clean envT+)]
         (wrap (<tag> envT+' unquantifiedT)))])
    ([#.UnivQ] [#.ExQ])
    ))