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

##  Copyright (c) Eduardo Julian. All rights reserved.
##  This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
##  If a copy of the MPL was not distributed with this file,
##  You can obtain one at http://mozilla.org/MPL/2.0/.

(;module:
  lux
  (lux (control monad)
       (data [text "Text/" Eq<Text>]
             text/format
             [number]
             (struct [list "List/" Monad<List> Fold<List>]
                     [dict])
             [bool]
             [product])
       [compiler #+ Monad<Lux>]
       (macro [ast]
              ["s" syntax #+ syntax: Syntax])
       [type]
       (type ["tc" check #+ Check Monad<Check>])
       ))

(def: (find-type-var id env)
  (-> Nat (Bindings Nat (Maybe Type)) (Lux Type))
  (case (list;find (|>. product;left (n.= id))
                   (get@ #;mappings env))
    (#;Some [_ (#;Some type)])
    (case type
      (#;VarT id')
      (find-type-var id' env)

      _
      (:: Monad<Lux> wrap type))

    (#;Some [_ #;None])
    (compiler;fail (format "Unbound type-var " (%n id)))

    #;None
    (compiler;fail (format "Unknown type-var " (%n id)))
    ))

(def: (resolve-type var-name)
  (-> Ident (Lux Type))
  (do Monad<Lux>
    [raw-type (compiler;find-type var-name)
     compiler compiler;get-compiler]
    (case raw-type
      (#;VarT id)
      (find-type-var id (get@ #;type-vars compiler))

      _
      (wrap raw-type))))

(def: (find-member-type idx sig-type)
  (-> Nat Type (Check Type))
  (case sig-type
    (#;NamedT _ sig-type')
    (find-member-type idx sig-type')

    (#;AppT func arg)
    (case (type;apply-type func arg)
      #;None
      (tc;fail (format "Can't apply type " (%type func) " to type " (%type arg)))

      (#;Some sig-type')
      (find-member-type idx sig-type'))

    (#;ProdT left right)
    (if (n.= +0 idx)
      (:: Monad<Check> wrap left)
      (find-member-type (n.dec idx) right))

    _
    (if (n.= +0 idx)
      (:: Monad<Check> wrap sig-type)
      (tc;fail (format "Can't find member type " (%n idx) " for " (%type sig-type))))))

(def: (find-member-name member)
  (-> Ident (Lux Ident))
  (case member
    ["" simple-name]
    (compiler;either (do Monad<Lux>
                       [member (compiler;normalize member)
                        _ (compiler;resolve-tag member)]
                       (wrap member))
                     (do Monad<Lux>
                       [this-module-name compiler;current-module-name
                        imp-mods (compiler;imported-modules this-module-name)
                        tag-lists (mapM @ compiler;tag-lists imp-mods)
                        #let [tag-lists (|> tag-lists List/join (List/map product;left) List/join)
                              candidates (list;filter (. (Text/= simple-name) product;right)
                                                      tag-lists)]]
                       (case candidates
                         #;Nil
                         (compiler;fail (format "Unknown tag: " (%ident member)))
                         
                         (#;Cons winner #;Nil)
                         (wrap winner)
                         
                         _
                         (compiler;fail (format "Too many candidate tags: " (%list %ident candidates))))))

    _
    (:: Monad<Lux> wrap member)))

(def: (resolve-member member)
  (-> Ident (Lux [Nat Type]))
  (do Monad<Lux>
    [member (find-member-name member)
     [idx tag-list sig-type] (compiler;resolve-tag member)]
    (wrap [idx sig-type])))

(def: (prepare-defs this-module-name defs)
  (-> Text (List [Text Def]) (List [Ident Type]))
  (|> defs
      (list;filter (lambda [[name [def-type def-anns def-value]]]
                     (compiler;struct? def-anns)))
      (List/map (lambda [[name [def-type def-anns def-value]]]
                  [[this-module-name name] def-type]))))

(def: local-env
  (Lux (List [Ident Type]))
  (do Monad<Lux>
    [local-batches compiler;locals
     #let [total-locals (List/fold (lambda [[name type] table]
                                     (dict;put~ name type table))
                                   (: (dict;Dict Text Type)
                                      (dict;new text;Hash<Text>))
                                   (List/join local-batches))]]
    (wrap (|> total-locals
              dict;entries
              (List/map (lambda [[name type]] [["" name] type]))))))

(def: local-structs
  (Lux (List [Ident Type]))
  (do Monad<Lux>
    [this-module-name compiler;current-module-name
     defs (compiler;defs this-module-name)]
    (wrap (prepare-defs this-module-name defs))))

(def: import-structs
  (Lux (List [Ident Type]))
  (do Monad<Lux>
    [this-module-name compiler;current-module-name
     imp-mods (compiler;imported-modules this-module-name)
     export-batches (mapM @ (lambda [imp-mod]
                              (do @
                                [exports (compiler;exports imp-mod)]
                                (wrap (prepare-defs imp-mod exports))))
                          imp-mods)]
    (wrap (List/join export-batches))))

(def: (apply-function-type func arg)
  (-> Type Type (Check Type))
  (case func
    (#;NamedT _ func')
    (apply-function-type func' arg)

    (#;UnivQ _)
    (do Monad<Check>
      [[id var] tc;create-var]
      (apply-function-type (default (undefined)
                             (type;apply-type func var))
                           arg))

    (#;LambdaT input output)
    (do Monad<Check>
      [_ (tc;check input arg)]
      (wrap output))

    _
    (tc;fail (format "Invalid function type: " (%type func)))))

(def: (concrete-type type)
  (-> Type (Check [(List Nat) Type]))
  (case type
    (#;UnivQ _)
    (do Monad<Check>
      [[id var] tc;create-var
       [ids final-output] (concrete-type (default (undefined)
                                           (type;apply-type type var)))]
      (wrap [(#;Cons id ids)
             final-output]))
    
    _
    (:: Monad<Check> wrap [(list) type])))

(def: (check-apply member-type input-types output-type)
  (-> Type (List Type) Type (Check []))
  (do Monad<Check>
    [member-type' (foldM Monad<Check>
                         (lambda [input member]
                           (apply-function-type member input))
                         member-type
                         input-types)]
    (tc;check output-type member-type')))

(def: compiler-type-context
  (Lux tc;Context)
  (lambda [compiler]
    (let [type-vars (get@ #;type-vars compiler)
          context (|> tc;fresh-context
                      (set@ #tc;var-counter (get@ #;counter type-vars))
                      (set@ #tc;bindings (dict;from-list number;Hash<Nat> (get@ #;mappings type-vars))))]
      (#;Right [compiler context]))))

(def: (test-provision context dep alts)
  (-> tc;Context Type (List [Ident Type]) (Lux (List Ident)))
  (do Monad<Lux>
    [compiler compiler;get-compiler]
    (case (|> alts
              (list;filter (lambda [[alt-name alt-type]]
                             (case (tc;run context
                                           (do Monad<Check>
                                             [_ (tc;check dep alt-type)]
                                             (wrap [])))
                               (#;Left error)
                               false

                               (#;Right _)
                               true)))
              (List/map product;left))
      #;Nil
      (compiler;fail (format "No candidates for provisioning: " (%type dep)))

      found
      (wrap found))))

(def: (provision compiler context dep)
  (-> Compiler tc;Context Type (Check Ident))
  (case (compiler;run compiler
                      ($_ compiler;either
                          (do Monad<Lux> [alts local-env] (test-provision context dep alts))
                          (do Monad<Lux> [alts local-structs] (test-provision context dep alts))
                          (do Monad<Lux> [alts import-structs] (test-provision context dep alts))))
    (#;Left error)
    (tc;fail error)

    (#;Right candidates)
    (case candidates
      #;Nil
      (tc;fail (format "No candidates for provisioning: " (%type dep)))

      (#;Cons winner #;Nil)
      (:: Monad<Check> wrap winner)

      _
      (tc;fail (format "Too many candidates for provisioning: " (%type dep) " --- " (%list %ident candidates))))
    ))

(def: (test-alternatives sig-type member-idx input-types output-type alts)
  (-> Type Nat (List Type) Type (List [Ident Type]) (Lux (List [Ident (List Ident)])))
  (do Monad<Lux>
    [compiler compiler;get-compiler
     context compiler-type-context]
    (case (|> alts
              (List/map (lambda [[alt-name alt-type]]
                          (case (tc;run context
                                        (do Monad<Check>
                                          [[tvars alt-type] (concrete-type alt-type)
                                           #let [[deps alt-type] (type;flatten-function alt-type)]
                                           _ (tc;check alt-type sig-type)
                                           member-type (find-member-type member-idx alt-type)
                                           _ (check-apply member-type input-types output-type)
                                           =tvars (mapM @ tc;deref tvars)
                                           context' tc;get-context
                                           =deps (mapM @ (provision compiler context') deps)]
                                          (wrap =deps)))
                            (#;Left error)
                            (list)

                            (#;Right =deps)
                            (list [alt-name =deps]))))
              List/join)
      #;Nil
      (compiler;fail (format "No alternatives."
                             "\n"
                             (|> alts
                                 (List/map product;left)
                                 (%list %ident))))

      found
      (wrap found))))

(def: (find-alternatives sig-type member-idx input-types output-type)
  (-> Type Nat (List Type) Type (Lux (List [Ident (List Ident)])))
  (let [test (test-alternatives sig-type member-idx input-types output-type)]
    ($_ compiler;either
        (do Monad<Lux> [alts local-env] (test alts))
        (do Monad<Lux> [alts local-structs] (test alts))
        (do Monad<Lux> [alts import-structs] (test alts)))))

(def: (var? input)
  (-> AST Bool)
  (case input
    [_ (#;SymbolS _)]
    true

    _
    false))

(def: (join-pair [l r])
  (All [a] (-> [a a] (List a)))
  (list l r))

(syntax: #export (::: [member s;symbol]
                      [args (s;alt (s;seq (s;some s;symbol) s;end)
                                   (s;seq (s;some s;any) s;end))])
  (case args
    (#;Left [args _])
    (do @
      [[member-idx sig-type] (resolve-member member)
       input-types (mapM @ resolve-type args)
       output-type compiler;expected-type
       chosen-ones (find-alternatives sig-type member-idx input-types output-type)]
      (case chosen-ones
        #;Nil
        (compiler;fail (format "No structure option could be found for member: " (%ident member)))

        (#;Cons [chosen deps] #;Nil)
        (let [chosen-inst (case deps
                            #;Nil
                            (ast;symbol chosen)

                            _
                            (` ((~ (ast;symbol chosen)) (~@ (List/map ast;symbol deps)))))]
          (wrap (list (` (:: (~ chosen-inst)
                             (~ (ast;local-symbol (product;right member)))
                             (~@ (List/map ast;symbol args)))))))

        _
        (compiler;fail (format "Too many options available: "
                               (|> chosen-ones
                                   (List/map (. %ident product;left))
                                   (text;join-with ", ")
                                   )))))

    (#;Right [args _])
    (do @
      [labels (seqM @ (list;repeat (list;size args)
                                   (compiler;gensym "")))
       #let [retry (` (let [(~@ (|> (list;zip2 labels args) (List/map join-pair) List/join))]
                        (;;::: (~ (ast;symbol member)) (~@ labels))))]]
      (wrap (list retry)))
    ))