path: root/stdlib/source/lux/lang/type/check.lux

(.module: {#.doc "Type-checking functionality."}
  lux
  (lux (control [functor #+ Functor]
                [apply #+ Apply]
                [monad #+ do Monad]
                ["ex" exception #+ exception:])
       (data [text "text/" Monoid<Text> Eq<Text>]
             [number "nat/" Codec<Text,Nat>]
             [maybe]
             [product]
             (coll [list]
                   (set ["set" unordered #+ Set]))
             ["e" error])
       (lang [type "type/" Eq<Type>])
       ))

(exception: #export (unknown-type-var {id Nat})
  (nat/encode id))

(exception: #export (unbound-type-var {id Nat})
  (nat/encode id))

(exception: #export (invalid-type-application {funcT Type} {argT Type})
  (type.to-text (#.Apply argT funcT)))

(exception: #export (cannot-rebind-var {id Nat} {type Type} {bound Type})
  (ex.report ["Var" (nat/encode id)]
             ["Wanted Type" (type.to-text type)]
             ["Current Type" (type.to-text bound)]))

(exception: #export (type-check-failed {expected Type} {actual Type})
  (ex.report ["Expected" (type.to-text expected)]
             ["Actual" (type.to-text actual)]))

(type: #export Var Nat)

(type: #export Assumption
  {#subsumption [Type Type]
   #verdict Bool})

(type: #export (Check a)
  (-> Type-Context (e.Error [Type-Context a])))

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

(struct: #export _ (Functor Check)
  (def: (map f fa)
    (function (_ context)
      (case (fa context)
        (#e.Error error)
        (#e.Error error)

        (#e.Success [context' output])
        (#e.Success [context' (f output)])
        ))))

(struct: #export _ (Apply Check)
  (def: functor Functor<Check>)

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

          (#e.Error error)
          (#e.Error error))

        (#e.Error error)
        (#e.Error error)
        )))
  )

(struct: #export _ (Monad Check)
  (def: functor Functor<Check>)

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

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

          (#e.Error error)
          (#e.Error error))

        (#e.Error error)
        (#e.Error error)
        )))
  )

(open: "check/" Monad<Check>)

(def: (var::get id plist)
  (-> Var Type-Vars (Maybe (Maybe Type)))
  (case plist
    #.Nil
    #.None

    (#.Cons [var-id var-type]
            plist')
    (if (n/= id var-id)
      (#.Some var-type)
      (var::get id plist'))
    ))

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

(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: (var::remove id plist)
  (-> Var Type-Vars Type-Vars)
  (case plist
    #.Nil
    #.Nil

    (#.Cons [var-id var-type]
            plist')
    (if (n/= id var-id)
      plist'
      (#.Cons [var-id var-type]
              (var::remove id plist')))
    ))

## [[Logic]]
(def: #export (run context proc)
  (All [a] (-> Type-Context (Check a) (e.Error a)))
  (case (proc context)
    (#e.Error error)
    (#e.Error error)

    (#e.Success [context' output])
    (#e.Success output)))

(def: #export (throw exception message)
  (All [e a] (-> (ex.Exception e) e (Check a)))
  (function (_ context)
    (ex.throw exception message)))

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

(do-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))
         (#e.Success [context <fail>])
         
         (#.Some (#.Some bound))
         (#e.Success [context <succeed>])

         #.None
         (ex.throw unknown-type-var id))))]

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

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

      (#.Some #.None)
      (ex.throw unbound-type-var id)

      #.None
      (ex.throw unknown-type-var id))))

(def: #export (write type id)
  (-> Type Var (Check Any))
  (function (_ context)
    (case (|> context (get@ #.var-bindings) (var::get id))
      (#.Some (#.Some bound))
      (ex.throw cannot-rebind-var [id type bound])
      
      (#.Some #.None)
      (#e.Success [(update@ #.var-bindings (var::put id (#.Some type)) context)
                   []])

      #.None
      (ex.throw unknown-type-var id))))

(def: (update type id)
  (-> Type Var (Check Any))
  (function (_ context)
    (case (|> context (get@ #.var-bindings) (var::get id))
      (#.Some _)
      (#e.Success [(update@ #.var-bindings (var::put id (#.Some type)) context)
                   []])
      
      #.None
      (ex.throw unknown-type-var id))))

(def: #export var
  (Check [Var Type])
  (function (_ context)
    (let [id (get@ #.var-counter context)]
      (#e.Success [(|> context
                       (update@ #.var-counter inc)
                       (update@ #.var-bindings (var::new id)))
                   [id (#.Var id)]]))))

(def: get-bindings
  (Check (List [Var (Maybe Type)]))
  (function (_ context)
    (#e.Success [context
                 (get@ #.var-bindings context)])))

(def: (set-bindings value)
  (-> (List [Var (Maybe Type)]) (Check Any))
  (function (_ context)
    (#e.Success [(set@ #.var-bindings value context)
                 []])))

(def: (apply-type! funcT argT)
  (-> Type Type (Check Type))
  (case funcT
    (#.Var func-id)
    (do Monad<Check>
      [?funcT' (read func-id)]
      (case ?funcT'
        #.None
        (throw invalid-type-application [funcT argT])

        (#.Some funcT')
        (apply-type! funcT' argT)))

    _
    (function (_ context)
      (case (type.apply (list argT) funcT)
        #.None
        (ex.throw invalid-type-application [funcT argT])

        (#.Some output)
        (#e.Success [context output])))))

(type: #export Ring (Set Var))

(def: empty-ring Ring (set.new number.Hash<Nat>))

(def: #export (ring id)
  (-> Var (Check Ring))
  (function (_ context)
    (loop [current id
           output (set.add id empty-ring)]
      (case (|> context (get@ #.var-bindings) (var::get current))
        (#.Some (#.Some type))
        (case type
          (#.Var post)
          (if (n/= id post)
            (#e.Success [context output])
            (recur post (set.add post output)))
          
          _
          (#e.Success [context empty-ring]))

        (#.Some #.None)
        (#e.Success [context output])
        
        #.None
        (ex.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)
      (#e.Success [context' output])
      (#e.Success [context' (#.Some output)])

      (#e.Error _)
      (#e.Success [context #.None]))))

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

(def: #export (assert message test)
  (-> Text Bool (Check Any))
  (function (_ context)
    (if test
      (#e.Success [context []])
      (#e.Error message))))

(def: (either left right)
  (All [a] (-> (Check a) (Check a) (Check a)))
  (function (_ context)
    (case (left context)
      (#e.Success [context' output])
      (#e.Success [context' output])

      (#e.Error _)
      (right context))))

(def: (assumed? [e a] assumptions)
  (-> [Type Type] (List Assumption) (Maybe Bool))
  (:: maybe.Monad<Maybe> map product.right
      (list.find (function (_ [[fe fa] status])
                   (and (type/= e fe)
                        (type/= a fa)))
                 assumptions)))

(def: (assume! ea status assumptions)
  (-> [Type Type] Bool (List Assumption) (List Assumption))
  (#.Cons [ea status] assumptions))

(def: (on id type then else)
  (All [a]
    (-> Var Type (Check a) (-> Type (Check a))
        (Check a)))
  ($_ either
      (do Monad<Check>
        [_ (write type id)]
        then)
      (do Monad<Check>
        [ring (ring id)
         _ (assert "" (n/> +1 (set.size ring)))
         _ (monad.map @ (update type) (set.to-list ring))]
        then)
      (do Monad<Check>
        [?bound (read id)]
        (else (maybe.default (#.Var id) ?bound)))))

(def: (link-2 left right)
  (-> Var Var (Check Any))
  (do Monad<Check>
    [_ (write (#.Var right) left)]
    (write (#.Var left) right)))

(def: (link-3 interpose to from)
  (-> Var Var Var (Check Any))
  (do Monad<Check>
    [_ (update (#.Var interpose) from)]
    (update (#.Var to) interpose)))

(def: (check-vars check' assumptions idE idA)
  (-> (-> Type Type (List Assumption) (Check (List Assumption)))
      (List Assumption)
      Var Var
      (Check (List Assumption)))
  (if (n/= idE idA)
    (check/wrap assumptions)
    (do Monad<Check>
      [ebound (attempt (peek idE))
       abound (attempt (peek idA))]
      (case [ebound abound]
        ## Link the 2 variables circularily
        [#.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' etype (#.Var idA) assumptions))

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

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

        [(#.Some etype) (#.Some atype)]
        (case [etype atype]
          [(#.Var targetE) (#.Var targetA)]
          (do @
            [ringE (ring idE)
             ringA (ring idA)]
            (if (:: set.Eq<Set> = 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))))
          
          [(#.Var targetE) _]
          (do @
            [ring (ring idE)
             _ (monad.map @ (update atype) (set.to-list ring))]
            (wrap assumptions))
          
          [_ (#.Var targetA)]
          (do @
            [ring (ring idA)
             _ (monad.map @ (update etype) (set.to-list ring))]
            (wrap assumptions))
          
          _
          (check' etype atype assumptions))))))

(def: (with-error-stack on-error check)
  (All [a] (-> (-> Any Text) (Check a) (Check a)))
  (function (_ context)
    (case (check context)
      (#e.Error error)
      (#e.Error (case error
                  ""
                  (on-error [])

                  _
                  ($_ text/compose
                      (on-error [])
                      "\n\n-----------------------------------------\n\n"
                      error)))

      output
      output)))

(def: (check-apply check' assumptions [eAT eFT] [aAT aFT])
  (-> (-> Type Type (List Assumption) (Check (List Assumption))) (List Assumption)
      [Type Type] [Type Type]
      (Check (List Assumption)))
  (case [eFT aFT]
    (^or [(#.UnivQ _ _) (#.Ex _)] [(#.UnivQ _ _) (#.Var _)])
    (do Monad<Check>
      [eFT' (apply-type! eFT eAT)]
      (check' eFT' (#.Apply aAT aFT) assumptions))

    (^or [(#.Ex _) (#.UnivQ _ _)] [(#.Var _) (#.UnivQ _ _)])
    (do Monad<Check>
      [aFT' (apply-type! aFT aAT)]
      (check' (#.Apply eAT eFT) aFT' assumptions))

    (^or [(#.Ex _) _] [_ (#.Ex _)])
    (do Monad<Check>
      [assumptions (check' eFT aFT assumptions)]
      (check' eAT aAT assumptions))

    [(#.Var id) _]
    (do Monad<Check>
      [?rFT (read id)]
      (case ?rFT
        (#.Some rFT)
        (check' (#.Apply eAT rFT) (#.Apply aAT aFT) assumptions)

        _
        (do Monad<Check>
          [assumptions (check' eFT aFT assumptions)
           e' (apply-type! aFT eAT)
           a' (apply-type! aFT aAT)]
          (check' e' a' assumptions))))

    [_ (#.Var id)]
    (do Monad<Check>
      [?rFT (read id)]
      (case ?rFT
        (#.Some rFT)
        (check' (#.Apply eAT eFT) (#.Apply aAT rFT) assumptions)

        _
        (do Monad<Check>
          [assumptions (check' eFT aFT assumptions)
           e' (apply-type! eFT eAT)
           a' (apply-type! eFT aAT)]
          (check' e' a' assumptions))))

    _
    (fail "")))

(def: #export (check' expected actual assumptions)
  {#.doc "Type-check to ensure that the 'expected' type subsumes the 'actual' type."}
  (-> Type Type (List Assumption) (Check (List Assumption)))
  (if (is? expected actual)
    (check/wrap assumptions)
    (with-error-stack
      (function (_ _) (ex.construct type-check-failed [expected actual]))
      (case [expected actual]
        [(#.Var idE) (#.Var idA)]
        (check-vars check' assumptions idE idA)
        
        [(#.Var id) _]
        (on id actual
            (check/wrap assumptions)
            (function (_ bound)
              (check' bound actual assumptions)))
        
        [_ (#.Var id)]
        (on id expected
            (check/wrap assumptions)
            (function (_ bound)
              (check' expected bound assumptions)))

        (^template [<fe> <fa>]
          [(#.Apply A1 <fe>) (#.Apply A2 <fa>)]
          (check-apply check' assumptions [A1 <fe>] [A2 <fa>]))
        ([F1 (#.Ex ex)]
         [(#.Ex ex) F2]
         [F1 (#.Var id)]
         [(#.Var id) F2])
        
        [(#.Apply A F) _]
        (let [fx-pair [expected actual]]
          (case (assumed? fx-pair assumptions)
            (#.Some ?)
            (if ?
              (check/wrap assumptions)
              (fail ""))

            #.None
            (do Monad<Check>
              [expected' (apply-type! F A)]
              (check' expected' actual (assume! fx-pair true assumptions)))))

        [_ (#.Apply A F)]
        (do Monad<Check>
          [actual' (apply-type! F A)]
          (check' expected actual' assumptions))

        (^template [<tag> <instancer>]
          [(<tag> _) _]
          (do Monad<Check>
            [[_ paramT] <instancer>
             expected' (apply-type! expected paramT)]
            (check' expected' actual assumptions)))
        ([#.UnivQ ..existential]
         [#.ExQ ..var])

        (^template [<tag> <instancer>]
          [_ (<tag> _)]
          (do Monad<Check>
            [[_ paramT] <instancer>
             actual' (apply-type! actual paramT)]
            (check' expected actual' assumptions)))
        ([#.UnivQ ..var]
         [#.ExQ ..existential])

        [(#.Primitive e-name e-params) (#.Primitive a-name a-params)]
        (if (and (text/= e-name a-name)
                 (n/= (list.size e-params)
                      (list.size a-params)))
          (do Monad<Check>
            [assumptions (monad.fold Monad<Check>
                                     (function (_ [e a] assumptions) (check' e a assumptions))
                                     assumptions
                                     (list.zip2 e-params a-params))]
            (check/wrap assumptions))
          (fail ""))

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

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

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

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

        _
        (fail "")))))

(def: #export (check expected actual)
  {#.doc "Type-check to ensure that the 'expected' type subsumes the 'actual' type."}
  (-> Type Type (Check Any))
  (do Monad<Check>
    [assumptions (check' expected actual (list))]
    (wrap [])))

(def: #export (checks? expected actual)
  {#.doc "A simple type-checking function that just returns a yes/no answer."}
  (-> Type Type Bool)
  (case (run fresh-context (check expected actual))
    (#e.Error error)
    false

    (#e.Success _)
    true))

(def: #export get-context
  (Check Type-Context)
  (function (_ context)
    (#e.Success [context context])))

(def: #export (clean inputT)
  (-> Type (Check Type))
  (case inputT
    (#.Primitive name paramsT+)
    (do Monad<Check>
      [paramsT+' (monad.map @ clean paramsT+)]
      (wrap (#.Primitive name paramsT+')))

    (^or (#.Parameter _) (#.Ex _) (#.Named _))
    (:: Monad<Check> wrap inputT)

    (^template [<tag>]
      (<tag> leftT rightT)
      (do Monad<Check>
        [leftT' (clean leftT)
         rightT' (clean rightT)]
        (wrap (<tag> leftT' rightT'))))
    ([#.Sum] [#.Product] [#.Function] [#.Apply])

    (#.Var id)
    (do Monad<Check>
      [?actualT (read id)]
      (case ?actualT
        (#.Some actualT)
        (clean actualT)

        _
        (wrap inputT)))

    (^template [<tag>]
      (<tag> envT+ unquantifiedT)
      (do Monad<Check>
        [envT+' (monad.map @ clean envT+)]
        (wrap (<tag> envT+' unquantifiedT))))
    ([#.UnivQ] [#.ExQ])
    ))