-- | Entry point for Hydra type inference, which is a variation on on Hindley-Milner

module Hydra.Inference.Inference (
  inferTermType,
  inferGraphTypes,
  inferredTypeOf,
  inferTypeScheme,
  inferTypeAndConstraints,
  withInferenceContext,
) where

import Hydra.Compute
import Hydra.Core
import Hydra.CoreEncoding
import Hydra.Graph
import Hydra.Lexical
import Hydra.Mantle
import Hydra.Annotations
import Hydra.Rewriting
import Hydra.Inference.Substitution
import Hydra.Unification
import Hydra.Inference.Rules
import Hydra.Tier1
import Hydra.Tier2
import Hydra.Tools.Sorting
import qualified Hydra.Dsl.Terms as Terms
import qualified Hydra.Dsl.Types as Types

import qualified Control.Monad as CM
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.Maybe as Y


annotateElements :: Graph -> [Element] -> Flow Graph [Element]
annotateElements :: Graph -> [Element] -> Flow Graph [Element]
annotateElements Graph
g [Element]
sortedEls = Flow Graph [Element] -> Flow Graph [Element]
forall {b}. Flow Graph b -> Flow Graph b
withInferenceContext (Flow Graph [Element] -> Flow Graph [Element])
-> Flow Graph [Element] -> Flow Graph [Element]
forall a b. (a -> b) -> a -> b
$ do
    [(Element, [TypeConstraint])]
iels' <- [Element]
-> [(Element, [TypeConstraint])]
-> Flow Graph [(Element, [TypeConstraint])]
annotate [Element]
sortedEls ([])
    let iels :: [Element]
iels = (Element, [TypeConstraint]) -> Element
forall a b. (a, b) -> a
fst ((Element, [TypeConstraint]) -> Element)
-> [(Element, [TypeConstraint])] -> [Element]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Element, [TypeConstraint])]
iels'
    let constraints :: [[TypeConstraint]]
constraints = (Element, [TypeConstraint]) -> [TypeConstraint]
forall a b. (a, b) -> b
snd ((Element, [TypeConstraint]) -> [TypeConstraint])
-> [(Element, [TypeConstraint])] -> [[TypeConstraint]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Element, [TypeConstraint])]
iels'

    -- Note: inference occurs over the entire graph at once,
    --       but unification and substitution occur within elements in isolation
    [Subst]
subst <- Flow Graph [Subst] -> Flow Graph [Subst]
forall {b}. Flow Graph b -> Flow Graph b
withSchemaContext (Flow Graph [Subst] -> Flow Graph [Subst])
-> Flow Graph [Subst] -> Flow Graph [Subst]
forall a b. (a -> b) -> a -> b
$ ([TypeConstraint] -> Flow Graph Subst)
-> [[TypeConstraint]] -> Flow Graph [Subst]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
CM.mapM [TypeConstraint] -> Flow Graph Subst
forall s. [TypeConstraint] -> Flow s Subst
solveConstraints [[TypeConstraint]]
constraints
    [Element] -> Flow Graph [Element]
forall a. a -> Flow Graph a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Element] -> Flow Graph [Element])
-> [Element] -> Flow Graph [Element]
forall a b. (a -> b) -> a -> b
$ (Subst -> Element -> Element) -> [Subst] -> [Element] -> [Element]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
L.zipWith Subst -> Element -> Element
rewriteElement [Subst]
subst [Element]
iels
  where
    -- Note: the following defaults to user-provided type annotations where provided.
    --       In the future, we should trust unification to perform this defaulting, and not override the inferred type.
    rewriteElement :: Subst -> Element -> Element
rewriteElement Subst
subst Element
el = Element
el { elementData = setTermType (Just typ) term1 }
      where
        term0 :: Term
term0 = Element -> Term
elementData Element
el
        term1 :: Term
term1 = (Type -> Type) -> Term -> Term
rewriteDataType (Subst -> Type -> Type
substituteInType Subst
subst) Term
term0
        typ :: Type
typ = Type -> Maybe Type -> Type
forall a. a -> Maybe a -> a
Y.fromMaybe (Term -> Type
termType Term
term1) (Maybe Type -> Type) -> Maybe Type -> Type
forall a b. (a -> b) -> a -> b
$ Term -> Maybe Type
getTermType Term
term0

    annotate :: [Element] -> [(Element, [TypeConstraint])] -> Flow Graph [(Element, [TypeConstraint])]
    annotate :: [Element]
-> [(Element, [TypeConstraint])]
-> Flow Graph [(Element, [TypeConstraint])]
annotate [Element]
original [(Element, [TypeConstraint])]
annotated = case [Element]
original of
      [] -> [(Element, [TypeConstraint])]
-> Flow Graph [(Element, [TypeConstraint])]
forall a. a -> Flow Graph a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(Element, [TypeConstraint])]
 -> Flow Graph [(Element, [TypeConstraint])])
-> [(Element, [TypeConstraint])]
-> Flow Graph [(Element, [TypeConstraint])]
forall a b. (a -> b) -> a -> b
$ [(Element, [TypeConstraint])] -> [(Element, [TypeConstraint])]
forall a. [a] -> [a]
L.reverse [(Element, [TypeConstraint])]
annotated
      (Element
el:[Element]
r) -> do
        (Inferred Element
iel Type
t [TypeConstraint]
c1) <- Element -> Flow Graph (Inferred Element)
inferElementType Element
el
        Name
-> TypeScheme
-> Flow Graph [(Element, [TypeConstraint])]
-> Flow Graph [(Element, [TypeConstraint])]
forall x. Name -> TypeScheme -> Flow Graph x -> Flow Graph x
withBinding (Element -> Name
elementName Element
el) (Type -> TypeScheme
monotype Type
t) (Flow Graph [(Element, [TypeConstraint])]
 -> Flow Graph [(Element, [TypeConstraint])])
-> Flow Graph [(Element, [TypeConstraint])]
-> Flow Graph [(Element, [TypeConstraint])]
forall a b. (a -> b) -> a -> b
$ [Element]
-> [(Element, [TypeConstraint])]
-> Flow Graph [(Element, [TypeConstraint])]
annotate [Element]
r ((Element
iel, [TypeConstraint]
c1)(Element, [TypeConstraint])
-> [(Element, [TypeConstraint])] -> [(Element, [TypeConstraint])]
forall a. a -> [a] -> [a]
:[(Element, [TypeConstraint])]
annotated)

inferTermType :: Term -> Flow Graph Term
inferTermType :: Term -> Flow Graph Term
inferTermType Term
term0 = do
  (Term
term1, TypeScheme
_) <- Term -> Flow Graph (Term, TypeScheme)
inferTypeAndConstraints Term
term0
  Term -> Flow Graph Term
forall a. a -> Flow Graph a
forall (m :: * -> *) a. Monad m => a -> m a
return Term
term1

inferElementType :: Element -> Flow Graph (Inferred Element)
inferElementType :: Element -> Flow Graph (Inferred Element)
inferElementType Element
el = String
-> Flow Graph (Inferred Element) -> Flow Graph (Inferred Element)
forall s a. String -> Flow s a -> Flow s a
withTrace (String
"infer type of " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
unName (Element -> Name
elementName Element
el)) (Flow Graph (Inferred Element) -> Flow Graph (Inferred Element))
-> Flow Graph (Inferred Element) -> Flow Graph (Inferred Element)
forall a b. (a -> b) -> a -> b
$ do
  (Inferred Term
iterm Type
t [TypeConstraint]
c) <- Term -> Flow Graph (Inferred Term)
infer (Term -> Flow Graph (Inferred Term))
-> Term -> Flow Graph (Inferred Term)
forall a b. (a -> b) -> a -> b
$ Element -> Term
elementData Element
el
  Inferred Element -> Flow Graph (Inferred Element)
forall a. a -> Flow Graph a
forall (m :: * -> *) a. Monad m => a -> m a
return (Element -> Type -> [TypeConstraint] -> Inferred Element
forall a. a -> Type -> [TypeConstraint] -> Inferred a
Inferred (Element
el {elementData = iterm}) Type
t [TypeConstraint]
c)

inferGraphTypes :: Flow Graph Graph
inferGraphTypes :: Flow Graph Graph
inferGraphTypes = Flow Graph Graph
forall s. Flow s s
getState Flow Graph Graph -> (Graph -> Flow Graph Graph) -> Flow Graph Graph
forall a b. Flow Graph a -> (a -> Flow Graph b) -> Flow Graph b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Graph -> Flow Graph Graph
annotateGraph
  where
    annotateGraph :: Graph -> Flow Graph Graph
annotateGraph Graph
g = String -> Flow Graph Graph -> Flow Graph Graph
forall s a. String -> Flow s a -> Flow s a
withTrace (String
"infer graph types") (Flow Graph Graph -> Flow Graph Graph)
-> Flow Graph Graph -> Flow Graph Graph
forall a b. (a -> b) -> a -> b
$ do
        [Element]
sorted <- Graph -> Flow Graph [Element]
sortGraphElements Graph
g
        [Element]
els <- Graph -> Flow Graph [Element]
sortGraphElements Graph
g Flow Graph [Element]
-> ([Element] -> Flow Graph [Element]) -> Flow Graph [Element]
forall a b. Flow Graph a -> (a -> Flow Graph b) -> Flow Graph b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Graph -> [Element] -> Flow Graph [Element]
annotateElements Graph
g
        Graph -> Flow Graph Graph
forall a. a -> Flow Graph a
forall (m :: * -> *) a. Monad m => a -> m a
return Graph
g {graphElements = M.fromList (toPair <$> els)}
      where
        toPair :: Element -> (Name, Element)
toPair Element
el = (Element -> Name
elementName Element
el, Element
el)

-- TODO: deprecated
inferredTypeOf :: Term -> Flow Graph Type
inferredTypeOf :: Term -> Flow Graph Type
inferredTypeOf Term
term = TypeScheme -> Type
typeSchemeType (TypeScheme -> Type) -> Flow Graph TypeScheme -> Flow Graph Type
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Flow Graph TypeScheme
inferTypeScheme Term
term

-- TODO: deprecated
-- | Solve for the top-level type of an expression in a given environment
inferTypeAndConstraints :: Term -> Flow Graph (Term, TypeScheme)
inferTypeAndConstraints :: Term -> Flow Graph (Term, TypeScheme)
inferTypeAndConstraints Term
term = String
-> Flow Graph (Term, TypeScheme) -> Flow Graph (Term, TypeScheme)
forall s a. String -> Flow s a -> Flow s a
withTrace (String
"infer type") (Flow Graph (Term, TypeScheme) -> Flow Graph (Term, TypeScheme))
-> Flow Graph (Term, TypeScheme) -> Flow Graph (Term, TypeScheme)
forall a b. (a -> b) -> a -> b
$ Flow Graph (Term, TypeScheme) -> Flow Graph (Term, TypeScheme)
forall {b}. Flow Graph b -> Flow Graph b
withInferenceContext (Flow Graph (Term, TypeScheme) -> Flow Graph (Term, TypeScheme))
-> Flow Graph (Term, TypeScheme) -> Flow Graph (Term, TypeScheme)
forall a b. (a -> b) -> a -> b
$ do
    (Inferred Term
iterm Type
_ [TypeConstraint]
constraints) <- Term -> Flow Graph (Inferred Term)
infer Term
term
    Subst
subst <- Flow Graph Subst -> Flow Graph Subst
forall {b}. Flow Graph b -> Flow Graph b
withSchemaContext (Flow Graph Subst -> Flow Graph Subst)
-> Flow Graph Subst -> Flow Graph Subst
forall a b. (a -> b) -> a -> b
$ [TypeConstraint] -> Flow Graph Subst
forall s. [TypeConstraint] -> Flow s Subst
solveConstraints [TypeConstraint]
constraints
    let term2 :: Term
term2 = (Type -> Type) -> Term -> Term
rewriteDataType (Subst -> Type -> Type
substituteInType Subst
subst) Term
iterm
--    let typ = Y.fromMaybe (termType term2) $ getTermType term
--    return (setTermType (Just typ) term2, closeOver $ termType term2)
    (Term, TypeScheme) -> Flow Graph (Term, TypeScheme)
forall a. a -> Flow Graph a
forall (m :: * -> *) a. Monad m => a -> m a
return (Term
term2, Type -> TypeScheme
closeOver (Type -> TypeScheme) -> Type -> TypeScheme
forall a b. (a -> b) -> a -> b
$ Term -> Type
termType Term
term2)
  where
    -- | Canonicalize and return the polymorphic top-level type.
    closeOver :: Type -> TypeScheme
closeOver = TypeScheme -> TypeScheme
normalizeScheme (TypeScheme -> TypeScheme)
-> (Type -> TypeScheme) -> Type -> TypeScheme
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name TypeScheme -> Type -> TypeScheme
generalize Map Name TypeScheme
forall k a. Map k a
M.empty (Type -> TypeScheme) -> (Type -> Type) -> Type -> TypeScheme
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Type
reduceType

-- TODO: deprecated
inferTypeScheme :: Term -> Flow Graph TypeScheme
inferTypeScheme :: Term -> Flow Graph TypeScheme
inferTypeScheme Term
term = (Term, TypeScheme) -> TypeScheme
forall a b. (a, b) -> b
snd ((Term, TypeScheme) -> TypeScheme)
-> Flow Graph (Term, TypeScheme) -> Flow Graph TypeScheme
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Flow Graph (Term, TypeScheme)
inferTypeAndConstraints Term
term

rewriteDataType :: (Type -> Type) -> Term -> Term
rewriteDataType :: (Type -> Type) -> Term -> Term
rewriteDataType Type -> Type
f = ((Term -> Term) -> Term -> Term) -> Term -> Term
rewriteTerm (Term -> Term) -> Term -> Term
forall {t}. (t -> Term) -> t -> Term
ff
  where
    ff :: (t -> Term) -> t -> Term
ff t -> Term
recurse t
term = case t -> Term
recurse t
term of
      TermTyped (TypedTerm Term
term1 Type
type1) -> TypedTerm -> Term
TermTyped (TypedTerm -> Term) -> TypedTerm -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Type -> TypedTerm
TypedTerm Term
term1 (Type -> Type
f Type
type1)
      Term
t -> Term
t

sortGraphElements :: Graph -> Flow Graph [Element]
sortGraphElements :: Graph -> Flow Graph [Element]
sortGraphElements Graph
g = do
    let annotated :: Set Name
annotated = [Name] -> Set Name
forall a. Ord a => [a] -> Set a
S.fromList ([Name] -> Set Name) -> [Name] -> Set Name
forall a b. (a -> b) -> a -> b
$ [Maybe Name] -> [Name]
forall a. [Maybe a] -> [a]
Y.catMaybes (Element -> Maybe Name
ifAnnotated (Element -> Maybe Name) -> [Element] -> [Maybe Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name Element -> [Element]
forall k a. Map k a -> [a]
M.elems Map Name Element
els)
    [(Name, [Name])]
adjList <- (Element -> Flow Graph (Name, [Name]))
-> [Element] -> Flow Graph [(Name, [Name])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
CM.mapM (Set Name -> Element -> Flow Graph (Name, [Name])
forall {m :: * -> *}.
Monad m =>
Set Name -> Element -> m (Name, [Name])
toAdj Set Name
annotated) ([Element] -> Flow Graph [(Name, [Name])])
-> [Element] -> Flow Graph [(Name, [Name])]
forall a b. (a -> b) -> a -> b
$ Map Name Element -> [Element]
forall k a. Map k a -> [a]
M.elems Map Name Element
els
    case [(Name, [Name])] -> Either [[Name]] [Name]
forall a. Ord a => [(a, [a])] -> Either [[a]] [a]
topologicalSort [(Name, [Name])]
adjList of
      Left [[Name]]
comps -> String -> Flow Graph [Element]
forall a. String -> Flow Graph a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Flow Graph [Element]) -> String -> Flow Graph [Element]
forall a b. (a -> b) -> a -> b
$ String
"cyclical dependency not resolved through annotations: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
L.intercalate String
", " (Name -> String
unName (Name -> String) -> [Name] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[Name]] -> [Name]
forall a. HasCallStack => [a] -> a
L.head [[Name]]
comps)
      Right [Name]
names -> [Element] -> Flow Graph [Element]
forall a. a -> Flow Graph a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Element] -> Flow Graph [Element])
-> [Element] -> Flow Graph [Element]
forall a b. (a -> b) -> a -> b
$ [Maybe Element] -> [Element]
forall a. [Maybe a] -> [a]
Y.catMaybes ((\Name
n -> Name -> Map Name Element -> Maybe Element
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
n Map Name Element
els) (Name -> Maybe Element) -> [Name] -> [Maybe Element]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
names)
  where
    els :: Map Name Element
els = Graph -> Map Name Element
graphElements Graph
g
    ifAnnotated :: Element -> Maybe Name
ifAnnotated Element
el = case (Term -> Maybe Type
getTermType (Term -> Maybe Type) -> Term -> Maybe Type
forall a b. (a -> b) -> a -> b
$ Element -> Term
elementData Element
el) of
      Maybe Type
Nothing -> Maybe Name
forall a. Maybe a
Nothing
      Just Type
_ -> Name -> Maybe Name
forall a. a -> Maybe a
Just (Name -> Maybe Name) -> Name -> Maybe Name
forall a b. (a -> b) -> a -> b
$ Element -> Name
elementName Element
el
    toAdj :: Set Name -> Element -> m (Name, [Name])
toAdj Set Name
annotated Element
el = do
        let deps :: [Name]
deps = (Name -> Bool) -> [Name] -> [Name]
forall a. (a -> Bool) -> [a] -> [a]
L.filter Name -> Bool
isNotAnnotated ([Name] -> [Name]) -> [Name] -> [Name]
forall a b. (a -> b) -> a -> b
$ (Name -> Bool) -> [Name] -> [Name]
forall a. (a -> Bool) -> [a] -> [a]
L.filter Name -> Bool
isElName ([Name] -> [Name]) -> [Name] -> [Name]
forall a b. (a -> b) -> a -> b
$ Set Name -> [Name]
forall a. Set a -> [a]
S.toList (Set Name -> [Name]) -> Set Name -> [Name]
forall a b. (a -> b) -> a -> b
$ Term -> Set Name
freeVariablesInTerm (Term -> Set Name) -> Term -> Set Name
forall a b. (a -> b) -> a -> b
$ Element -> Term
elementData Element
el

        (Name, [Name]) -> m (Name, [Name])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Element -> Name
elementName Element
el, [Name]
deps)
      where
        -- Ignore free variables which are not valid element references
        isElName :: Name -> Bool
isElName Name
name = Name -> Map Name Element -> Bool
forall k a. Ord k => k -> Map k a -> Bool
M.member Name
name Map Name Element
els
        -- No need for an inference dependency on an element which is already annotated with a type
        isNotAnnotated :: Name -> Bool
isNotAnnotated Name
name = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Name
name Set Name
annotated

withInferenceContext :: Flow Graph b -> Flow Graph b
withInferenceContext Flow Graph b
flow = do
    Graph
g <- Flow Graph Graph
forall s. Flow s s
getState
    let env :: Map Name TypeScheme
env = Graph -> Map Name TypeScheme
initialEnv Graph
g
    Graph -> Flow Graph b -> Flow Graph b
forall s1 a s2. s1 -> Flow s1 a -> Flow s2 a
withState (Graph
g {graphTypes = env}) Flow Graph b
flow
  where
    initialEnv :: Graph -> Map Name TypeScheme
initialEnv Graph
g = [(Name, TypeScheme)] -> Map Name TypeScheme
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, TypeScheme)] -> Map Name TypeScheme)
-> [(Name, TypeScheme)] -> Map Name TypeScheme
forall a b. (a -> b) -> a -> b
$ [Maybe (Name, TypeScheme)] -> [(Name, TypeScheme)]
forall a. [Maybe a] -> [a]
Y.catMaybes (Element -> Maybe (Name, TypeScheme)
toPair (Element -> Maybe (Name, TypeScheme))
-> [Element] -> [Maybe (Name, TypeScheme)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map Name Element -> [Element]
forall k a. Map k a -> [a]
M.elems (Map Name Element -> [Element]) -> Map Name Element -> [Element]
forall a b. (a -> b) -> a -> b
$ Graph -> Map Name Element
graphElements Graph
g))
      where
        toPair :: Element -> Maybe (Name, TypeScheme)
toPair Element
el = (\Type
t -> (Element -> Name
elementName Element
el, Type -> TypeScheme
monotype Type
t)) (Type -> (Name, TypeScheme))
-> Maybe Type -> Maybe (Name, TypeScheme)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Term -> Maybe Type
getTermType (Term -> Maybe Type) -> Term -> Maybe Type
forall a b. (a -> b) -> a -> b
$ Element -> Term
elementData Element
el)