path: root/stdlib/source/lux/macro/poly.lux

(;module:
  [lux #- function]
  (lux (control monad
                [eq])
       (data [text]
             text/format
             (coll [list "List/" Fold<List> Monad<List>]
                   [dict #+ Dict])
             [number]
             [product]
             [bool]
             [char]
             [maybe]
             [ident "Ident/" Eq<Ident>])
       [macro #+ Monad<Lux> with-gensyms]
       (macro [code]
              ["s" syntax #+ syntax: Syntax]
              (syntax [common]))
       [type]
       ))

## [Types]
(type: #export (Matcher a)
  (-> Type (Lux a)))

(type: #export Env (Dict Nat [Type Code]))

## [Combinators]
(do-template [<combinator> <name> <type>]
  [(def: #export <combinator>
     (Matcher Unit)
     (;function [:type:]
       (case (type;un-name :type:)
         <type>
         (:: macro;Monad<Lux> wrap [])

         _
         (macro;fail (format "Not " <name> " type: " (%type :type:))))))]

  [void "Void" #;Void]
  [unit "Unit" #;Unit]
  )

(do-template [<combinator> <name>]
  [(def: #export <combinator>
     (Matcher Unit)
     (;function [:type:]
       (case (type;un-alias :type:)
         (#;Named ["lux" <name>] _)
         (:: macro;Monad<Lux> wrap [])

         _
         (macro;fail (format "Not " <name> " type: " (%type :type:))))))]

  [bool "Bool"]
  [nat  "Nat"]
  [int  "Int"]
  [deg  "Deg"]
  [real "Real"]
  [char "Char"]
  [text "Text"]
  )

(def: #export primitive
  (Matcher Type)
  (;function [:type:]
    (with-expansions
      [<primitives> (do-template [<parser> <type>]
                      [(do Monad<Lux>
                         [_ (<parser> :type:)]
                         (wrap <type>))]

                      [void Void]
                      [unit Unit]
                      [bool Bool]
                      [nat  Nat]
                      [int  Int]
                      [deg  Deg]
                      [real Real]
                      [char Char]
                      [text Text])]
      ($_ macro;either
          <primitives>))))

(syntax: ($Code$ ast)
  (wrap (;list (code;text (code;to-text ast)))))

(do-template [<single> <multi> <flattener> <tag>]
  [(def: #export <single>
     (Matcher [Type Type])
     (;function [:type:]
       (case (type;un-name :type:)
         (<tag> :left: :right:)
         (:: macro;Monad<Lux> wrap [:left: :right:])

         _
         (macro;fail (format "Not a " ($Code$ <tag>) " type: " (%type :type:))))))

   (def: #export <multi>
     (Matcher (List Type))
     (;function [:type:]
       (let [members (<flattener> (type;un-name :type:))]
         (if (n.> +1 (list;size members))
           (:: macro;Monad<Lux> wrap members)
           (macro;fail (format "Not a " ($Code$ <tag>) " type: " (%type :type:)))))))]

  [sum    sum+    type;flatten-variant #;Sum]
  [prod   prod+   type;flatten-tuple   #;Product]
  )

(def: #export func
  (Matcher [Type Type])
  (;function [:type:]
    (case (type;un-name :type:)
      (#;Function :left: :right:)
      (:: macro;Monad<Lux> wrap [:left: :right:])

      _
      (macro;fail (format "Not a Function type: " (%type :type:))))))

(def: #export func+
  (Matcher [(List Type) Type])
  (;function [:type:]
    (let [[ins out] (type;flatten-function (type;un-name :type:))]
      (if (n.> +0 (list;size ins))
        (:: macro;Monad<Lux> wrap [ins out])
        (macro;fail (format "Not a Function type: " (%type :type:)))))))

(def: #export tagged
  (Matcher [(List Ident) Type])
  (;function [:type:]
    (case (type;un-alias :type:)
      (#;Named type-name :def:)
      (do macro;Monad<Lux>
        [tags (macro;tags-of type-name)]
        (wrap [tags :def:]))

      _
      (macro;fail (format "Unnamed types cannot have tags: " (%type :type:))))))

(def: #export polymorphic
  (Matcher [(List Code) Type])
  (;function [:type:]
    (loop [:type: (type;un-name :type:)]
      (case :type:
        (#;UnivQ _ :type:')
        (do macro;Monad<Lux>
          [[g!tail :type:''] (recur :type:')
           g!head (macro;gensym "type-var")]
          (wrap [(list& g!head g!tail)
                 :type:'']))

        _
        (:: macro;Monad<Lux> wrap [(;list) :type:])))))

(do-template [<combinator> <sub-comb> <build>]
  [(def: #export <combinator>
     (Matcher [(List Code) (List [Ident Type])])
     (;function [:type:]
       (do macro;Monad<Lux>
         [[tags :type:] (tagged :type:)
          _ (macro;assert "Records and variants must have tags."
                          (n.> +0 (list;size tags)))
          [vars :type:] (polymorphic :type:)
          members (<sub-comb> :type:)
          #let [num-tags (list;size tags)
                [init-tags last-tag] (list;split (n.dec num-tags) tags)
                [init-types last-types] (list;split (n.dec num-tags) members)]]
         (wrap [vars (list;concat (;list (list;zip2 init-tags init-types)
                                         (;list [(default (undefined)
                                                   (list;head last-tag))
                                                 (<build> last-types)])))]))))]

  [variant sum+  type;variant]
  [record  prod+ type;tuple]
  )

(def: #export tuple
  (Matcher [(List Code) (List Type)])
  (;function [:type:]
    (do macro;Monad<Lux>
      [[vars :type:] (polymorphic :type:)
       members (prod+ :type:)]
      (wrap [vars members]))))

(def: #export function
  (Matcher [(List Code) [(List Type) Type]])
  (;function [:type:]
    (do macro;Monad<Lux>
      [[vars :type:] (polymorphic :type:)
       ins+out (func+ :type:)]
      (wrap [vars ins+out]))))

(def: #export apply
  (Matcher [Type (List Type)])
  (;function [:type:]
    (do macro;Monad<Lux>
      [#let [[:func: :args:] (loop [:type: (type;un-name :type:)]
                               (case :type:
                                 (#;App :func: :arg:)
                                 (let [[:func:' :args:] (recur :func:)]
                                   [:func:' (list& :arg: :args:)])

                                 _
                                 [:type: (;list)]))]]
      (case :args:
        #;Nil
        (macro;fail "Not a type application.")

        _
        (wrap [:func: (list;reverse :args:)])))))

(def: #export (apply-1 name)
  (-> Ident (Matcher Type))
  (;function [:type:]
    (case (type;un-name :type:)
      (^multi (#;App :quant: :arg:)
              [(type;un-alias :quant:) (#;Named actual _)]
              (Ident/= name actual))
      (:: macro;Monad<Lux> wrap :arg:)

      _
      (macro;fail (format "Not " (%ident name) " type: " (%type :type:))))))

(def: #export (apply-2 name)
  (-> Ident (Matcher [Type Type]))
  (;function [:type:]
    (case (type;un-name :type:)
      (^multi (#;App (#;App :quant: :arg0:) :arg1:)
              [(type;un-alias :quant:) (#;Named actual _)]
              (Ident/= name actual))
      (:: macro;Monad<Lux> wrap [:arg0: :arg1:])

      _
      (macro;fail (format "Not " (%ident name) " type: " (%type :type:))))))

(def: (adjusted-idx env idx)
  (-> Env Nat Nat)
  (let [env-level (n./ +2 (dict;size env))
        bound-level (n./ +2 idx)
        bound-idx (n.% +2 idx)]
    (|> env-level n.dec (n.- bound-level) (n.* +2) (n.+ bound-idx))))

(def: #export (bound env)
  (-> Env (Matcher Code))
  (;function [:type:]
    (case :type:
      (#;Bound idx)
      (case (dict;get (adjusted-idx env idx) env)
        (#;Some [poly-type poly-ast])
        (:: macro;Monad<Lux> wrap poly-ast)

        #;None
        (macro;fail (format "Unknown bound type: " (%type :type:))))

      _
      (macro;fail (format "Not a bound type: " (%type :type:))))))

(def: #export (recur env)
  (-> Env (Matcher Code))
  (;function [:type:]
    (do Monad<Lux>
      [[t-func t-args] (apply :type:)]
      (case t-func
        (^multi (#;Bound t-func-idx)
                (n.= +0 (adjusted-idx env t-func-idx))
                [(do maybe;Monad<Maybe>
                   [=func (dict;get +0 env)
                    =args (mapM @ (;function [t-arg]
                                    (case t-arg
                                      (#;Bound idx)
                                      (dict;get (adjusted-idx env idx) env)

                                      _
                                      #;None))
                                t-args)]
                   (wrap (` ((~ (product;right =func)) (~@ (List/map product;right =args))))))
                 (#;Some call)])
        (wrap call)

        _
        (macro;fail (format "Type is not a recursive instance: " (%type :type:))))
      )))

(def: #export (var env var-id)
  (-> Env Nat (Matcher Unit))
  (;function [:type:]
    (case :type:
      (^multi (#;Bound idx)
              (n.= var-id (adjusted-idx env idx)))
      (:: macro;Monad<Lux> wrap [])

      _
      (macro;fail (format "Not a bound type: " (%type :type:))))))

## [Syntax]
(def: #export (extend-env [funcT funcA] type-vars env)
  (-> [Type Code] (List [Type Code]) Env Env)
  (case type-vars
    #;Nil
    env
    
    (#;Cons [varT varA] type-vars')
    (let [current-size (dict;size env)]
      (|> env
          (dict;put current-size [funcT funcA])
          (dict;put (n.inc current-size) [varT varA])
          (extend-env [(#;App funcT varT) (` (#;App (~ funcA) (~ varA)))]
                      type-vars')
          ))))

(syntax: #export (poly: [_ex-lev common;export-level]
                   [[name env inputs] (s;form ($_ s;seq
                                                  s;local-symbol
                                                  s;local-symbol
                                                  (s;many s;local-symbol)))]
                   body)
  (with-gensyms [g!body]
    (let [g!inputs (List/map (|>. [""] code;symbol) inputs)
          g!name (code;symbol ["" name])
          g!env (code;symbol ["" env])]
      (wrap (;list (` (syntax: (~@ (common;gen-export-level _ex-lev)) ((~ g!name) (~@ (List/map (;function [g!input] (` [(~ g!input) s;symbol]))
                                                                                                g!inputs)))
                        (do Monad<Lux>
                          [(~@ (List/join (List/map (;function [g!input] (;list g!input (` (macro;find-type-def (~ g!input)))))
                                                    g!inputs)))
                           (~' #let) [(~ g!env) (: Env (dict;new number;Hash<Nat>))]
                           (~ g!body) (: (Lux Code)
                                         (loop [(~ g!env) (~ g!env)
                                                (~@ (List/join (List/map (;function [g!input] (;list g!input g!input))
                                                                         g!inputs)))]
                                           (let [(~ g!name) (~' recur)]
                                             (~ body))))]
                          ((~' wrap) (;list (~ g!body)))))))))))

(def: (common-poly-name? poly-func)
  (-> Text Bool)
  (text;contains? "?" poly-func))

(def: (derivation-name poly args)
  (-> Text (List Text) (Maybe Text))
  (if (common-poly-name? poly)
    (#;Some (List/fold (text;replace-once "?") poly args))
    #;None))

(syntax: #export (derived: [_ex-lev common;export-level]
                   [?name (s;opt s;local-symbol)]
                   [[poly-func poly-args] (s;form (s;seq s;symbol (s;many s;symbol)))]
                   [?custom-impl (s;opt s;any)])
  (do @
    [poly-args (mapM @ macro;normalize poly-args)
     name (case ?name
            (#;Some name)
            (wrap name)

            (^multi #;None
                    [(derivation-name (product;right poly-func) (List/map product;right poly-args))
                     (#;Some derived-name)])
            (wrap derived-name)

            _
            (macro;fail "derived: was given no explicit name, and cannot generate one from given information."))
     #let [impl (case ?custom-impl
                  (#;Some custom-impl)
                  custom-impl

                  #;None
                  (` ((~ (code;symbol poly-func)) (~@ (List/map code;symbol poly-args)))))]]
    (wrap (;list (` (def: (~@ (common;gen-export-level _ex-lev))
                      (~ (code;symbol ["" name]))
                      {#;struct? true}
                      (~ impl)))))))

## [Derivers]
(def: (to-ast env type)
  (-> Env Type Code)
  (case type
    (#;Host name params)
    (` (#;Host (~ (code;text name))
               (list (~@ (List/map (to-ast env) params)))))

    (^template [<tag>]
      <tag>
      (` <tag>))
    ([#;Void] [#;Unit])

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

    (#;Bound idx)
    (let [idx (adjusted-idx env idx)]
      (if (n.= +0 idx)
        (|> (dict;get idx env) (default (undefined)) product;left (to-ast env))
        (` (;$ (~ (code;nat (n.dec idx)))))))

    (^template [<tag>]
      (<tag> left right)
      (` (<tag> (~ (to-ast env left))
                (~ (to-ast env right)))))
    ([#;Function] [#;App])

    (^template [<tag> <macro> <flattener>]
      (<tag> left right)
      (` (<macro> (~@ (List/map (to-ast env) (<flattener> type))))))
    ([#;Sum  | type;flatten-variant]
     [#;Product & type;flatten-tuple])

    (#;Named name sub-type)
    (code;symbol name)

    (^template [<tag>]
      (<tag> scope body)
      (` (<tag> (list (~@ (List/map (to-ast env) scope)))
                (~ (to-ast env body)))))
    ([#;UnivQ] [#;ExQ])
    ))

(def: #export (gen-type env converter type-fun tvars type)
  (-> Env (-> Code Code) Code (List Code) Type Code)
  (let [type' (to-ast env type)]
    (case tvars
      #;Nil
      (converter type')

      _
      (` (All (~ type-fun) [(~@ tvars)]
           (-> (~@ (List/map converter tvars))
               (~ (converter (` ((~ type') (~@ tvars)))))))))))

(def: #export (type-var-indices num-vars)
  (-> Nat (List Type))
  (|> num-vars list;indices (List/map (|>. #;Bound))))