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

(;module: {#;doc "Type-checking functionality.

                  Very useful for writing advanced macros."}
  lux
  (lux (control functor
                applicative
                monad)
       (data [text "Text/" Monoid<Text> Eq<Text>]
             text/format
             maybe
             [product]
             (coll [list])
             [error #- fail])
       [type "Type/" Eq<Type>]
       ))

(type: #export Fixed (List [[Type Type] Bool]))

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

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

(struct: #export _ (Functor Check)
  (def: (map f fa)
    (function [context]
      (case (fa context)
        (#;Left error)
        (#;Left error)

        (#;Right [context' output])
        (#;Right [context' (f output)])
        ))))

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

  (def: (wrap x)
    (function [context]
      (#;Right [context x])))

  (def: (apply ff fa)
    (function [context]
      (case (ff context)
        (#;Right [context' f])
        (case (fa context')
          (#;Right [context'' a])
          (#;Right [context'' (f a)])

          (#;Left error)
          (#;Left error))

        (#;Left error)
        (#;Left error)
        )))
  )

(struct: #export _ (Monad Check)
  (def: applicative Applicative<Check>)

  (def: (join ffa)
    (function [context]
      (case (ffa context)
        (#;Right [context' fa])
        (case (fa context')
          (#;Right [context'' a])
          (#;Right [context'' a])

          (#;Left error)
          (#;Left error))

        (#;Left error)
        (#;Left error)
        )))
  )

(open Monad<Check> "Check/")

(def: (var::get id plist)
  (-> Nat 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::put id value plist)
  (-> Nat (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)
  (-> Nat 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) (Error a)))
  (case (proc context)
    (#;Left error)
    (#;Left error)

    (#;Right [context' output])
    (#;Right output)))

(def: (apply-type! t-func t-arg)
  (-> Type Type (Check Type))
  (function [context]
    (case (type;apply-type t-func t-arg)
      #;None
      (#;Left (format "Invalid type application: " (%type t-func) " on " (%type t-arg)))

      (#;Some output)
      (#;Right [context output]))))

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

(def: #export (bound? id)
  (-> Nat (Check Bool))
  (function [context]
    (case (|> context (get@ #;var-bindings) (var::get id))
      (#;Some (#;Some _))
      (#;Right [context true])

      (#;Some #;None)
      (#;Right [context false])
      
      #;None
      (#;Left (format "Unknown type-var: " (%n id))))))

(def: #export (read-var id)
  (-> Nat (Check Type))
  (function [context]
    (case (|> context (get@ #;var-bindings) (var::get id))
      (#;Some (#;Some type))
      (#;Right [context type])

      (#;Some #;None)
      (#;Left (format "Unbound type-var: " (%n id)))
      
      #;None
      (#;Left (format "Unknown type-var: " (%n id))))))

(def: #export (write-var id type)
  (-> Nat Type (Check Unit))
  (function [context]
    (case (|> context (get@ #;var-bindings) (var::get id))
      (#;Some (#;Some bound))
      (#;Left (format "Cannot rebind type-var: " (%n id) " | Current type: " (%type bound)))
      
      (#;Some #;None)
      (#;Right [(update@ #;var-bindings (var::put id (#;Some type)) context)
                []])

      #;None
      (#;Left (format "Unknown type-var: " (%n id))))))

(def: (rewrite-var id type)
  (-> Nat Type (Check Unit))
  (function [context]
    (case (|> context (get@ #;var-bindings) (var::get id))
      (#;Some _)
      (#;Right [(update@ #;var-bindings (var::put id (#;Some type)) context)
                []])
      
      #;None
      (#;Left (format "Unknown type-var: " (%n id))))))

(def: #export (clear-var id)
  (-> Nat (Check Unit))
  (function [context]
    (case (|> context (get@ #;var-bindings) (var::get id))
      (#;Some _)
      (#;Right [(update@ #;var-bindings (var::put id #;None) context)
                []])
      
      #;None
      (#;Left (format "Unknown type-var: " (%n id))))))

(def: #export (clean t-id type)
  (-> Nat Type (Check Type))
  (case type
    (#;Var id)
    (if (n.= t-id id)
      (do Monad<Check>
        [? (bound? id)]
        (if ?
          (read-var id)
          (wrap type)))
      (do Monad<Check>
        [? (bound? id)]
        (if ?
          (do Monad<Check>
            [=type (read-var id)
             ==type (clean t-id =type)]
            (case ==type
              (#;Var =id)
              (if (n.= t-id =id)
                (do Monad<Check>
                  [_ (clear-var id)]
                  (wrap type))
                (do Monad<Check>
                  [_ (rewrite-var id ==type)]
                  (wrap type)))

              _
              (do Monad<Check>
                [_ (rewrite-var id ==type)]
                (wrap type))))
          (wrap type))))

    (#;Host name params)
    (do Monad<Check>
      [=params (mapM @ (clean t-id) params)]
      (wrap (#;Host name =params)))
    
    (^template [<tag>]
      (<tag> left right)
      (do Monad<Check>
        [=left (clean t-id left)
         =right (clean t-id right)]
        (wrap (<tag> =left =right))))
    ([#;Function]
     [#;App]
     [#;Product]
     [#;Sum])

    (^template [<tag>]
      (<tag> env body)
      (do Monad<Check>
        [=env (mapM @ (clean t-id) env)
         =body (clean t-id body)] ## TODO: DO NOT CLEAN THE BODY
        (wrap (<tag> =env =body))))
    ([#;UnivQ]
     [#;ExQ])
    
    _
    (:: Monad<Check> wrap type)
    ))

(def: #export create-var
  (Check [Nat Type])
  (function [context]
    (let [id (get@ #;var-counter context)]
      (#;Right [(|> context
                    (update@ #;var-counter n.inc)
                    (update@ #;var-bindings (var::put id #;None)))
                [id (#;Var id)]]))))

(def: get-bindings
  (Check (List [Nat (Maybe Type)]))
  (function [context]
    (#;Right [context
              (get@ #;var-bindings context)])))

(def: (set-bindings value)
  (-> (List [Nat (Maybe Type)]) (Check Unit))
  (function [context]
    (#;Right [(set@ #;var-bindings value context)
              []])))

(def: #export (delete-var id)
  (-> Nat (Check Unit))
  (do Monad<Check>
    [? (bound? id)
     _ (if ?
         (wrap [])
         (do Monad<Check>
           [[ex-id ex] existential]
           (write-var id ex)))
     bindings get-bindings
     bindings' (mapM @
                     (function [(^@ binding [b-id b-type])]
                       (if (n.= id b-id)
                         (wrap binding)
                         (case b-type
                           #;None
                           (wrap binding)

                           (#;Some b-type')
                           (case b-type'
                             (#;Var t-id)
                             (if (n.= id t-id)
                               (wrap [b-id #;None])
                               (wrap binding))

                             _
                             (do Monad<Check>
                               [b-type'' (clean id b-type')]
                               (wrap [b-id (#;Some b-type'')])))
                           )))
                     bindings)]
    (set-bindings (var::remove id bindings'))))

(def: #export (with-var k)
  (All [a] (-> (-> [Nat Type] (Check a)) (Check a)))
  (do Monad<Check>
    [[id var] create-var
     output (k [id var])
     _ (delete-var id)]
    (wrap output)))

(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)
      (#;Right [context' output])
      (#;Right [context' (#;Some output)])

      (#;Left _)
      (#;Right [context #;None]))))

(def: #export (fail message)
  (All [a] (-> Text (Check a)))
  (function [context]
    (#;Left message)))

(def: (fail-check expected actual)
  (All [a] (-> Type Type (Check a)))
  (fail (format "Expected: " (%type expected) "\n\n"
                "Actual:   " (%type actual))))

(def: (either left right)
  (All [a] (-> (Check a) (Check a) (Check a)))
  (function [context]
    (case (left context)
      (#;Right [context' output])
      (#;Right [context' output])

      (#;Left _)
      (right context))))

(def: (fx-get [e a] fixed)
  (-> [Type Type] Fixed (Maybe Bool))
  (:: Monad<Maybe> map product;right
      (list;find (function [[[fe fa] status]]
                   (and (Type/= e fe)
                        (Type/= a fa)))
                 fixed)))

(def: (fx-put ea status fixed)
  (-> [Type Type] Bool Fixed Fixed)
  (#;Cons [ea status] fixed))

(def: (on-var id type then else)
  (All [a]
    (-> Nat Type (Check a) (-> Type (Check a))
        (Check a)))
  (either (do Monad<Check>
            [_ (write-var id type)]
            then)
          (do Monad<Check>
            [bound (read-var id)]
            (else bound))))

(def: #export (check' expected actual fixed)
  {#;doc "Type-check to ensure that the 'expected' type subsumes the 'actual' type."}
  (-> Type Type Fixed (Check Fixed))
  (if (is expected actual)
    (Check/wrap fixed)
    (case [expected actual]
      [(#;Var e-id) (#;Var a-id)]
      (if (n.= e-id a-id)
        (Check/wrap fixed)
        (do Monad<Check>
          [ebound (attempt (read-var e-id))
           abound (attempt (read-var a-id))]
          (case [ebound abound]
            [#;None #;None]
            (do @
              [_ (write-var e-id actual)]
              (wrap fixed))
            
            [(#;Some etype) #;None]
            (check' etype actual fixed)

            [#;None (#;Some atype)]
            (check' expected atype fixed)

            [(#;Some etype) (#;Some atype)]
            (check' etype atype fixed))))
      
      [(#;Var id) _]
      (on-var id actual (Check/wrap fixed)
              (function [bound]
                (check' bound actual fixed)))
      
      [_ (#;Var id)]
      (on-var id expected (Check/wrap fixed)
              (function [bound]
                (check' expected bound fixed)))

      [(#;App (#;Ex eid) eA) (#;App (#;Ex aid) aA)]
      (if (n.= eid aid)
        (check' eA aA fixed)
        (fail-check expected actual))

      [(#;App (#;Var id) A1) (#;App F2 A2)]
      (either (do Monad<Check>
                [F1 (read-var id)]
                (check' (#;App F1 A1) actual fixed))
              (do Monad<Check>
                [fixed (check' (#;Var id) F2 fixed)
                 e' (apply-type! F2 A1)
                 a' (apply-type! F2 A2)]
                (check' e' a' fixed)))
      
      [(#;App F1 A1) (#;App (#;Var id) A2)]
      (either (do Monad<Check>
                [F2 (read-var id)]
                (check' expected (#;App F2 A2) fixed))
              (do Monad<Check>
                [fixed (check' F1 (#;Var id) fixed)
                 e' (apply-type! F1 A1)
                 a' (apply-type! F1 A2)]
                (check' e' a' fixed)))

      [(#;App F A) _]
      (let [fx-pair [expected actual]]
        (case (fx-get fx-pair fixed)
          (#;Some ?)
          (if ?
            (Check/wrap fixed)
            (fail-check expected actual))

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

      [_ (#;App F A)]
      (do Monad<Check>
        [actual' (apply-type! F A)]
        (check' expected actual' fixed))

      [(#;UnivQ _) _]
      (do Monad<Check>
        [[ex-id ex] existential
         expected' (apply-type! expected ex)]
        (check' expected' actual fixed))

      [_ (#;UnivQ _)]
      (with-var
        (function [[var-id var]]
          (do Monad<Check>
            [actual' (apply-type! actual var)
             fixed (check' expected actual' fixed)
             _ (clean var-id expected)]
            (Check/wrap fixed))))

      [(#;ExQ e!env e!def) _]
      (with-var
        (function [[var-id var]]
          (do Monad<Check>
            [expected' (apply-type! expected var)
             fixed (check' expected' actual fixed)
             _ (clean var-id actual)]
            (Check/wrap fixed))))

      [_ (#;ExQ a!env a!def)]
      (do Monad<Check>
        [[ex-id ex] existential
         actual' (apply-type! actual ex)]
        (check' expected actual' fixed))

      [(#;Host e-name e-params) (#;Host a-name a-params)]
      (if (and (Text/= e-name a-name)
               (n.= (list;size e-params)
                    (list;size a-params)))
        (do Monad<Check>
          [fixed (foldM Monad<Check>
                        (function [[e a] fixed] (check' e a fixed))
                        fixed
                        (list;zip2 e-params a-params))]
          (Check/wrap fixed))
        (fail-check expected actual))

      (^template [<unit> <append>]
        [<unit> <unit>]
        (Check/wrap fixed)
        
        [(<append> eL eR) (<append> aL aR)]
        (do Monad<Check>
          [fixed (check' eL aL fixed)]
          (check' eR aR fixed)))
      ([#;Void #;Sum]
       [#;Unit #;Product])
      
      [(#;Function eI eO) (#;Function aI aO)]
      (do Monad<Check>
        [fixed (check' aI eI fixed)]
        (check' eO aO fixed))

      [(#;Ex e!id) (#;Ex a!id)]
      (if (n.= e!id a!id)
        (Check/wrap fixed)
        (fail-check expected actual))

      [(#;Named _ ?etype) _]
      (check' ?etype actual fixed)

      [_ (#;Named _ ?atype)]
      (check' expected ?atype fixed)

      _
      (fail-check expected actual))))

(def: #export (check expected actual)
  {#;doc "Type-check to ensure that the 'expected' type subsumes the 'actual' type."}
  (-> Type Type (Check Unit))
  (do Monad<Check>
    [fixed (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))
    (#;Left error)
    false

    (#;Right _)
    true))

(def: #export get-context
  (Check Type-Context)
  (function [context]
    (#;Right [context context])))