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

(.module:
  [library
   [lux #*
    [abstract
     ["." monad (#+ do)]
     ["." equivalence]]
    [control
     ["." maybe]
     ["." try]
     ["<>" parser
      ["<.>" code (#+ Parser)]]]
    [data
     ["." product]
     ["." text ("#\." equivalence)
      ["%" format (#+ format)]]
     [collection
      ["." list ("#\." monad mix)]
      ["." dictionary (#+ Dictionary)]]]
    ["." macro
     ["." code]
     [syntax (#+ syntax:)]]
    [math
     ["." number
      ["n" nat]]]
    ["." meta
     ["." annotation]]
    ["." type
     ["." check (#+ Check)]]]])

(def: (type_var id env)
  (-> Nat Type_Context (Meta Type))
  (case (list.example (|>> product.left (n.= id))
                      (value@ #.var_bindings env))
    (#.Some [_ (#.Some type)])
    (case type
      (#.Var id')
      (type_var id' env)

      _
      (\ meta.monad in type))

    (#.Some [_ #.None])
    (meta.failure (format "Unbound type-var " (%.nat id)))

    #.None
    (meta.failure (format "Unknown type-var " (%.nat id)))
    ))

(def: (implicit_type var_name)
  (-> Name (Meta Type))
  (do meta.monad
    [raw_type (meta.type var_name)
     compiler meta.compiler_state]
    (case raw_type
      (#.Var id)
      (type_var id (value@ #.type_context compiler))

      _
      (in raw_type))))

(def: (member_type idx sig_type)
  (-> Nat Type (Check Type))
  (case sig_type
    (#.Named _ sig_type')
    (member_type idx sig_type')

    (#.Apply arg func)
    (case (type.applied (list arg) func)
      #.None
      (check.failure (format "Cannot apply type " (%.type func) " to type " (%.type arg)))

      (#.Some sig_type')
      (member_type idx sig_type'))

    (#.Product left right)
    (if (n.= 0 idx)
      (\ check.monad in left)
      (member_type (-- idx) right))

    _
    (if (n.= 0 idx)
      (\ check.monad in sig_type)
      (check.failure (format "Cannot find member type " (%.nat idx) " for " (%.type sig_type))))))

(def: (member_name member)
  (-> Name (Meta Name))
  (case member
    ["" simple_name]
    (meta.either (do meta.monad
                   [member (meta.normal member)
                    _ (meta.tag member)]
                   (in member))
                 (do {! meta.monad}
                   [this_module_name meta.current_module_name
                    imp_mods (meta.imported_modules this_module_name)
                    tag_lists (monad.map ! meta.tag_lists imp_mods)
                    .let [tag_lists (|> tag_lists list\join (list\map product.left) list\join)
                          candidates (list.only (|>> product.right (text\= simple_name))
                                                tag_lists)]]
                   (case candidates
                     #.End
                     (meta.failure (format "Unknown tag: " (%.name member)))
                     
                     (#.Item winner #.End)
                     (in winner)
                     
                     _
                     (meta.failure (format "Too many candidate tags: " (%.list %.name candidates))))))

    _
    (\ meta.monad in member)))

(def: (implicit_member member)
  (-> Name (Meta [Nat Type]))
  (do meta.monad
    [member (member_name member)
     [idx tag_list sig_type] (meta.tag member)]
    (in [idx sig_type])))

(def: (available_definitions source_module target_module constants aggregate)
  (-> Text Text (List [Text Definition]) (-> (List [Name Type]) (List [Name Type])))
  (list\mix (function (_ [name [exported? def_type def_anns def_value]] aggregate)
              (if (and (annotation.implementation? def_anns)
                       (or (text\= target_module source_module)
                           exported?))
                (#.Item [[source_module name] def_type] aggregate)
                aggregate))
            aggregate
            constants))

(def: local_env
  (Meta (List [Name Type]))
  (do meta.monad
    [local_batches meta.locals
     .let [total_locals (list\mix (function (_ [name type] table)
                                    (try.else table (dictionary.has' name type table)))
                                  (: (Dictionary Text Type)
                                     (dictionary.empty text.hash))
                                  (list\join local_batches))]]
    (in (|> total_locals
            dictionary.entries
            (list\map (function (_ [name type]) [["" name] type]))))))

(def: local_structs
  (Meta (List [Name Type]))
  (do {! meta.monad}
    [this_module_name meta.current_module_name
     definitions (meta.definitions this_module_name)]
    (in (available_definitions this_module_name this_module_name definitions #.End))))

(def: imported_structs
  (Meta (List [Name Type]))
  (do {! meta.monad}
    [this_module_name meta.current_module_name
     imported_modules (meta.imported_modules this_module_name)
     accessible_definitions (monad.map ! meta.definitions imported_modules)]
    (in (list\mix (function (_ [imported_module definitions] tail)
                    (available_definitions imported_module this_module_name definitions tail))
                  #.End
                  (list.zipped/2 imported_modules accessible_definitions)))))

(def: (on_argument arg func)
  (-> Type Type (Check Type))
  (case func
    (#.Named _ func')
    (on_argument arg func')

    (#.UnivQ _)
    (do check.monad
      [[id var] check.var]
      (|> func
          (type.applied (list var))
          maybe.trusted
          (on_argument arg)))

    (#.Function input output)
    (do check.monad
      [_ (check.check input arg)]
      (in output))

    _
    (check.failure (format "Invalid function type: " (%.type func)))))

(def: (concrete_type type)
  (-> Type (Check [(List Nat) Type]))
  (case type
    (#.UnivQ _)
    (do check.monad
      [[id var] check.var
       [ids final_output] (concrete_type (maybe.trusted (type.applied (list var) type)))]
      (in [(#.Item id ids)
           final_output]))
    
    _
    (\ check.monad in [(list) type])))

(def: (ensure_function_application! member_type input_types expected_output)
  (-> Type (List Type) Type (Check []))
  (do check.monad
    [actual_output (monad.mix check.monad ..on_argument member_type input_types)]
    (check.check expected_output actual_output)))

(type: #rec Instance
  {#constructor Name
   #dependencies (List Instance)})

(def: (candidate_provision provision context dep alts)
  (-> (-> Lux Type_Context Type (Check Instance))
      Type_Context Type (List [Name Type])
      (Meta (List Instance)))
  (do meta.monad
    [compiler meta.compiler_state]
    (case (|> alts
              (list\map (function (_ [alt_name alt_type])
                          (case (check.result context
                                              (do {! check.monad}
                                                [[tvars alt_type] (concrete_type alt_type)
                                                 .let [[deps alt_type] (type.flat_function alt_type)]
                                                 _ (check.check dep alt_type)
                                                 context' check.context
                                                 =deps (monad.map ! (provision compiler context') deps)]
                                                (in =deps)))
                            (#.Left error)
                            (list)

                            (#.Right =deps)
                            (list [alt_name =deps]))))
              list\join)
      #.End
      (meta.failure (format "No candidates for provisioning: " (%.type dep)))

      found
      (in found))))

(def: (provision compiler context dep)
  (-> Lux Type_Context Type (Check Instance))
  (case (meta.result compiler
                     ($_ meta.either
                         (do meta.monad [alts ..local_env] (..candidate_provision provision context dep alts))
                         (do meta.monad [alts ..local_structs] (..candidate_provision provision context dep alts))
                         (do meta.monad [alts ..imported_structs] (..candidate_provision provision context dep alts))))
    (#.Left error)
    (check.failure error)

    (#.Right candidates)
    (case candidates
      #.End
      (check.failure (format "No candidates for provisioning: " (%.type dep)))

      (#.Item winner #.End)
      (\ check.monad in winner)

      _
      (check.failure (format "Too many candidates for provisioning: " (%.type dep) " --- " (%.list (|>> product.left %.name) candidates))))
    ))

(def: (candidate_alternatives sig_type member_idx input_types output_type alts)
  (-> Type Nat (List Type) Type (List [Name Type]) (Meta (List Instance)))
  (do meta.monad
    [compiler meta.compiler_state
     context meta.type_context]
    (case (|> alts
              (list\map (function (_ [alt_name alt_type])
                          (case (<| (check.result context)
                                    (do {! check.monad}
                                      [[tvars alt_type] (concrete_type alt_type)
                                       .let [[deps alt_type] (type.flat_function alt_type)]
                                       _ (check.check alt_type sig_type)
                                       member_type (member_type member_idx alt_type)
                                       _ (ensure_function_application! member_type input_types output_type)
                                       context' check.context
                                       =deps (monad.map ! (provision compiler context') deps)]
                                      (in =deps)))
                            (#.Left error)
                            (list)

                            (#.Right =deps)
                            (list [alt_name =deps]))))
              list\join)
      #.End
      (meta.failure (format "No alternatives for " (%.type (type.function input_types output_type))))

      found
      (in found))))

(def: (alternatives sig_type member_idx input_types output_type)
  (-> Type Nat (List Type) Type (Meta (List Instance)))
  (let [test (candidate_alternatives sig_type member_idx input_types output_type)]
    ($_ meta.either
        (do meta.monad [alts ..local_env] (test alts))
        (do meta.monad [alts ..local_structs] (test alts))
        (do meta.monad [alts ..imported_structs] (test alts)))))

(def: (var? input)
  (-> Code Bit)
  (case input
    [_ (#.Identifier _)]
    #1

    _
    #0))

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

(def: (instance$ [constructor dependencies])
  (-> Instance Code)
  (case dependencies
    #.End
    (code.identifier constructor)

    _
    (` ((~ (code.identifier constructor)) (~+ (list\map instance$ dependencies))))))

(syntax: .public (\\ [member <code>.identifier
                      args (<>.or (<>.and (<>.some <code>.identifier) <code>.end!)
                                  (<>.and (<>.some <code>.any) <code>.end!))])
  {#.doc (example "Automatic implementation selection (for type-class style polymorphism)."
                  "This feature layers type-class style polymorphism on top of Lux's signatures and implementations."
                  "When calling a polymorphic function, or using a polymorphic constant,"
                  "this macro will check the types of the arguments, and the expected type for the whole expression"
                  "and it will search in the local scope, the module's scope and the imports' scope"
                  "in order to find suitable implementations to satisfy those requirements."
                  "If a single alternative is found, that one will be used automatically."
                  "If no alternative is found, or if more than one alternative is found (ambiguity)"
                  "a compile-time error will be raised, to alert the user."
                  "Examples:"
                  "Nat equivalence"
                  (\ number.equivalence = x y)
                  (\\ = x y)
                  "Can optionally add the prefix of the module where the signature was defined."
                  (\\ equivalence.= x y)
                  "(List Nat) equivalence"
                  (\\ =
                      (list.indices 10)
                      (list.indices 10))
                  "(Functor List) map"
                  (\\ map ++ (list.indices 10))
                  "Caveat emptor: You need to make sure to import the module of any implementation you want to use."
                  "Otherwise, this macro will not find it.")}
  (case args
    (#.Left [args _])
    (do {! meta.monad}
      [[member_idx sig_type] (..implicit_member member)
       input_types (monad.map ! ..implicit_type args)
       output_type meta.expected_type
       chosen_ones (alternatives sig_type member_idx input_types output_type)]
      (case chosen_ones
        #.End
        (meta.failure (format "No implementation could be found for member: " (%.name member)))

        (#.Item chosen #.End)
        (in (list (` (\ (~ (instance$ chosen))
                        (~ (code.local_identifier (product.right member)))
                        (~+ (list\map code.identifier args))))))

        _
        (meta.failure (format "Too many implementations available: "
                              (|> chosen_ones
                                  (list\map (|>> product.left %.name))
                                  (text.interposed ", "))
                              " --- for type: " (%.type sig_type)))))

    (#.Right [args _])
    (do {! meta.monad}
      [labels (|> (macro.identifier "") (list.repeated (list.size args)) (monad.all !))]
      (in (list (` (let [(~+ (|> args (list.zipped/2 labels) (list\map ..pair_list) list\join))]
                     (..\\ (~ (code.identifier member)) (~+ labels)))))))
    ))

(def: (implicit_bindings amount)
  (-> Nat (Meta (List Code)))
  (|> (macro.identifier "g!implicit")
      (list.repeated amount)
      (monad.all meta.monad)))

(def: implicits
  (Parser (List Code))
  (<code>.tuple (<>.many <code>.any)))

(syntax: .public (with [implementations ..implicits
                        body <code>.any])
  {#.doc (example "Establish lexical bindings for implementations that will be prioritized over non-lexically-bound implementations."
                  (with [n.addition]
                    (n.= (\ n.addition compose left right)
                         (\\ compose left right))))}
  (do meta.monad
    [g!implicit+ (implicit_bindings (list.size implementations))]
    (in (list (` (let [(~+ (|> (list.zipped/2 g!implicit+ implementations)
                               (list\map (function (_ [g!implicit implementation])
                                           (list g!implicit implementation)))
                               list\join))]
                   (~ body)))))))

(syntax: .public (implicit: [implementations ..implicits])
  {#.doc (example "Establish local definitions for implementations that will be prioritized over foreign definitions."
                  (implicit: [n.multiplication])
                  
                  (n.= (\ n.multiplication compose left right)
                       (\\ compose left right)))}
  (do meta.monad
    [g!implicit+ (implicit_bindings (list.size implementations))]
    (in (|> (list.zipped/2 g!implicit+ implementations)
            (list\map (function (_ [g!implicit implementation])
                        (` (def: .private (~ g!implicit)
                             {#.implementation? #1}
                             (~ implementation)))))))))