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

(.module: {#.doc "Basic functionality for working with types."}
  [lux (#- function)
   [abstract
    [equivalence (#+ Equivalence)]
    [monad (#+ Monad do)]]
   [control
    ["." function]
    ["." exception (#+ exception:)]
    ["p" parser
     ["s" code (#+ Parser)]]]
   [data
    ["." text ("#@." monoid equivalence)]
    ["." name ("#@." equivalence codec)]
    [number
     ["n" nat ("#@." decimal)]]
    ["." maybe]
    [collection
     ["." array]
     ["." list ("#@." functor monoid fold)]]]
   ["." macro
    ["." code]
    [syntax (#+ syntax:)]]])

(template [<name> <tag>]
  [(def: #export (<name> type)
     (-> Type [Nat Type])
     (loop [num-args 0
            type type]
       (case type
         (<tag> env sub-type)
         (recur (inc num-args) sub-type)

         _
         [num-args type])))]

  [flatten-univ-q #.UnivQ]
  [flatten-ex-q   #.ExQ]
  )

(def: #export (flatten-function type)
  (-> Type [(List Type) Type])
  (case type
    (#.Function in out')
    (let [[ins out] (flatten-function out')]
      [(list& in ins) out])

    _
    [(list) type]))

(def: #export (flatten-application type)
  (-> Type [Type (List Type)])
  (case type
    (#.Apply arg func')
    (let [[func args] (flatten-application func')]
      [func (list@compose args (list arg))])

    _
    [type (list)]))

(template [<name> <tag>]
  [(def: #export (<name> type)
     (-> Type (List Type))
     (case type
       (<tag> left right)
       (list& left (<name> right))

       _
       (list type)))]

  [flatten-variant #.Sum]
  [flatten-tuple   #.Product]
  )

(def: #export (to-text type)
  (-> Type Text)
  (case type
    (#.Primitive name params)
    ($_ text@compose
        "(primitive "
        (text.enclose' text.double-quote name)
        (|> params
            (list@map (|>> to-text (text@compose " ")))
            (list@fold (function.flip text@compose) ""))
        ")")

    (^template [<tag> <open> <close> <flatten>]
      (<tag> _)
      ($_ text@compose <open>
          (|> (<flatten> type)
              (list@map to-text)
              list.reverse
              (list.interpose " ")
              (list@fold text@compose ""))
          <close>))
    ([#.Sum  "(| " ")" flatten-variant]
     [#.Product "["   "]" flatten-tuple])

    (#.Function input output)
    (let [[ins out] (flatten-function type)]
      ($_ text@compose  "(-> "
          (|> ins
              (list@map to-text)
              list.reverse
              (list.interpose " ")
              (list@fold text@compose ""))
          " " (to-text out) ")"))

    (#.Parameter idx)
    (n@encode idx)

    (#.Var id)
    ($_ text@compose "⌈v:" (n@encode id) "⌋")

    (#.Ex id)
    ($_ text@compose "⟨e:" (n@encode id) "⟩")

    (#.Apply param fun)
    (let [[type-func type-args] (flatten-application type)]
      ($_ text@compose  "(" (to-text type-func) " " (|> type-args (list@map to-text) list.reverse (list.interpose " ") (list@fold text@compose "")) ")"))

    (^template [<tag> <desc>]
      (<tag> env body)
      ($_ text@compose "(" <desc> " {" (|> env (list@map to-text) (text.join-with " ")) "} " (to-text body) ")"))
    ([#.UnivQ "All"]
     [#.ExQ "Ex"])

    (#.Named [module name] type)
    ($_ text@compose module "." name)
    ))

(def: (beta-reduce env type)
  (-> (List Type) Type Type)
  (case type
    (#.Primitive name params)
    (#.Primitive name (list@map (beta-reduce env) params))
    
    (^template [<tag>]
      (<tag> left right)
      (<tag> (beta-reduce env left) (beta-reduce env right)))
    ([#.Sum]      [#.Product]
     [#.Function] [#.Apply])
    
    (^template [<tag>]
      (<tag> old-env def)
      (case old-env
        #.Nil
        (<tag> env def)

        _
        (<tag> (list@map (beta-reduce env) old-env) def)))
    ([#.UnivQ]
     [#.ExQ])
    
    (#.Parameter idx)
    (maybe.default (error! ($_ text@compose
                               "Unknown type parameter" text.new-line
                               "      Index: " (n@encode idx) text.new-line
                               "Environment: " (|> env
                                                   list.enumerate
                                                   (list@map (.function (_ [index type])
                                                               ($_ text@compose
                                                                   (n@encode index)
                                                                   " " (..to-text type))))
                                                   (text.join-with (text@compose text.new-line "             ")))))
                   (list.nth idx env))
    
    _
    type
    ))

(structure: #export equivalence (Equivalence Type)
  (def: (= x y)
    (or (is? x y)
        (case [x y]
          [(#.Primitive xname xparams) (#.Primitive yname yparams)]
          (and (text@= xname yname)
               (n.= (list.size yparams) (list.size xparams))
               (list@fold (.function (_ [x y] prev) (and prev (= x y)))
                          #1
                          (list.zip2 xparams yparams)))

          (^template [<tag>]
            [(<tag> xid) (<tag> yid)]
            (n.= yid xid))
          ([#.Var] [#.Ex] [#.Parameter])

          (^or [(#.Function xleft xright) (#.Function yleft yright)]
               [(#.Apply xleft xright) (#.Apply yleft yright)])
          (and (= xleft yleft)
               (= xright yright))

          [(#.Named xname xtype) (#.Named yname ytype)]
          (and (name@= xname yname)
               (= xtype ytype))

          (^template [<tag>]
            [(<tag> xL xR) (<tag> yL yR)]
            (and (= xL yL) (= xR yR)))
          ([#.Sum] [#.Product])
          
          (^or [(#.UnivQ xenv xbody) (#.UnivQ yenv ybody)]
               [(#.ExQ xenv xbody) (#.ExQ yenv ybody)])
          (and (n.= (list.size yenv) (list.size xenv))
               (= xbody ybody)
               (list@fold (.function (_ [x y] prev) (and prev (= x y)))
                          #1
                          (list.zip2 xenv yenv)))

          _
          #0
          ))))

(def: #export (apply params func)
  (-> (List Type) Type (Maybe Type))
  (case params
    #.Nil
    (#.Some func)

    (#.Cons param params')
    (case func
      (^template [<tag>]
        (<tag> env body)
        (|> body
            (beta-reduce (list& func param env))
            (apply params')))
      ([#.UnivQ] [#.ExQ])

      (#.Apply A F)
      (apply (list& A params) F)

      (#.Named name unnamed)
      (apply params unnamed)
      
      _
      #.None)))

(def: #export (to-code type)
  (-> Type Code)
  (case type
    (#.Primitive name params)
    (` (#.Primitive (~ (code.text name))
                    (.list (~+ (list@map to-code params)))))

    (^template [<tag>]
      (<tag> idx)
      (` (<tag> (~ (code.nat idx)))))
    ([#.Var] [#.Ex] [#.Parameter])

    (^template [<tag>]
      (<tag> left right)
      (` (<tag> (~ (to-code left))
                (~ (to-code right)))))
    ([#.Sum] [#.Product] [#.Function] [#.Apply])

    (#.Named name sub-type)
    (code.identifier name)

    (^template [<tag>]
      (<tag> env body)
      (` (<tag> (.list (~+ (list@map to-code env)))
                (~ (to-code body)))))
    ([#.UnivQ] [#.ExQ])
    ))

(def: #export (un-alias type)
  (-> Type Type)
  (case type
    (#.Named _ (#.Named name type'))
    (un-alias (#.Named name type'))

    _
    type))

(def: #export (un-name type)
  (-> Type Type)
  (case type
    (#.Named name type')
    (un-name type')

    _
    type))

(template [<name> <base> <ctor>]
  [(def: #export (<name> types)
     (-> (List Type) Type)
     (case types
       #.Nil
       <base>

       (#.Cons type #.Nil)
       type

       (#.Cons type types')
       (<ctor> type (<name> types'))))]

  [variant Nothing #.Sum]
  [tuple   Any     #.Product]
  )

(def: #export (function inputs output)
  (-> (List Type) Type Type)
  (case inputs
    #.Nil
    output

    (#.Cons input inputs')
    (#.Function input (function inputs' output))))

(def: #export (application params quant)
  (-> (List Type) Type Type)
  (case params
    #.Nil
    quant

    (#.Cons param params')
    (application params' (#.Apply param quant))))

(template [<name> <tag>]
  [(def: #export (<name> size body)
     (-> Nat Type Type)
     (case size
       0 body
       _  (|> body (<name> (dec size)) (<tag> (list)))))]

  [univ-q #.UnivQ]
  [ex-q   #.ExQ]
  )

(def: #export (quantified? type)
  (-> Type Bit)
  (case type
    (#.Named [module name] _type)
    (quantified? _type)

    (#.Apply A F)
    (maybe.default #0
                   (do maybe.monad
                     [applied (apply (list A) F)]
                     (wrap (quantified? applied))))
    
    (^or (#.UnivQ _) (#.ExQ _))
    #1

    _
    #0))

(def: #export (array depth elem-type)
  (-> Nat Type Type)
  (case depth
    0 elem-type
    _ (|> elem-type (array (dec depth)) (list) (#.Primitive array.type-name))))

(syntax: #export (:log! {input (p.or s.identifier
                                     s.any)})
  (case input
    (#.Left valueN)
    (do macro.monad
      [cursor macro.cursor
       valueT (macro.find-type valueN)
       #let [_ (log! ($_ text@compose
                         (name@encode (name-of ..:log!)) " @ " (.cursor-description cursor) text.new-line
                         "Value: " (name@encode valueN) text.new-line
                         " Type: " (..to-text valueT) text.new-line))]]
      (wrap (list (code.identifier valueN))))

    (#.Right valueC)
    (macro.with-gensyms [g!value]
      (wrap (list (` (.let [(~ g!value) (~ valueC)]
                       (..:log! (~ g!value)))))))))

(def: type-parameters
  (Parser (List Text))
  (s.tuple (p.some s.local-identifier)))

(syntax: #export (:cast {type-vars type-parameters}
                        input
                        output
                        {value (p.maybe s.any)})
  (let [casterC (` (: (All [(~+ (list@map code.local-identifier type-vars))]
                        (-> (~ input) (~ output)))
                      (|>> :assume)))]
    (case value
      #.None
      (wrap (list casterC))
      
      (#.Some value)
      (wrap (list (` ((~ casterC) (~ value))))))))

(type: Typed
  {#type Code
   #expression Code})

(def: typed
  (Parser Typed)
  (s.record (p.and s.any s.any)))

## TODO: Make sure the generated code always gets optimized away.
(syntax: #export (:share {type-vars type-parameters}
                         {exemplar typed}
                         {computation typed})
  (macro.with-gensyms [g!_]
    (let [shareC (` (: (All [(~+ (list@map code.local-identifier type-vars))]
                         (-> (~ (get@ #type exemplar))
                             (~ (get@ #type computation))))
                       (.function ((~ g!_) (~ g!_))
                         (~ (get@ #expression computation)))))]
      (wrap (list (` ((~ shareC) (~ (get@ #expression exemplar)))))))))

(syntax: #export (:by-example {type-vars type-parameters}
                              {exemplar typed}
                              {extraction s.any})
  (wrap (list (` (:of ((~! :share)
                       [(~+ (list@map code.local-identifier type-vars))]
                       {(~ (get@ #type exemplar))
                        (~ (get@ #expression exemplar))}
                       {(~ extraction)
                        (:assume [])}))))))

(exception: #export (hole-type {location Cursor} {type Type})
  (exception.report
   ["Location" (.cursor-description location)]
   ["Type" (..to-text type)]))

(syntax: #export (:hole)
  (do macro.monad
    [cursor macro.cursor
     expectedT macro.expected-type]
    (macro.fail (exception.construct ..hole-type [cursor expectedT]))))