{-# LANGUAGE LambdaCase #-}

module Telomare.TypeChecker where

import Control.Applicative
import Control.Comonad.Cofree (Cofree ((:<)))
import Control.Lens.Plated (transform)
import Control.Monad.Except
import Control.Monad.State (State)
import qualified Control.Monad.State as State
import Data.Bifunctor (second)
import qualified Data.DList as DList
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Debug.Trace
import PrettyPrint
import Telomare (FragExpr (..), FragExprF (..), FragExprUR (..), FragIndex (..),
                 LocTag (..), PartialType (..), RecursionSimulationPieces (..),
                 Term3 (..), rootFrag)

debug :: Bool
debug :: Bool
debug = Bool
False

debugTrace :: String -> a -> a
debugTrace :: forall a. String -> a -> a
debugTrace String
s a
x = if Bool
debug then forall a. String -> a -> a
trace String
s a
x else a
x

newtype DebugTypeMap = DebugTypeMap (Map Int PartialType)

instance Show DebugTypeMap where
  show :: DebugTypeMap -> String
show (DebugTypeMap Map Int PartialType
x) = (String
"typeMap:\n" forall a. [a] -> [a] -> [a]
++) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((forall a. [a] -> [a] -> [a]
++ String
"\n") forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second PartialType -> PrettyPartialType
PrettyPartialType) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
Map.toAscList Map Int PartialType
x

data TypeCheckError
  = UnboundType Int
  | InconsistentTypes PartialType PartialType
  | RecursiveType Int
  deriving (TypeCheckError -> TypeCheckError -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TypeCheckError -> TypeCheckError -> Bool
$c/= :: TypeCheckError -> TypeCheckError -> Bool
== :: TypeCheckError -> TypeCheckError -> Bool
$c== :: TypeCheckError -> TypeCheckError -> Bool
Eq, Eq TypeCheckError
TypeCheckError -> TypeCheckError -> Bool
TypeCheckError -> TypeCheckError -> Ordering
TypeCheckError -> TypeCheckError -> TypeCheckError
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: TypeCheckError -> TypeCheckError -> TypeCheckError
$cmin :: TypeCheckError -> TypeCheckError -> TypeCheckError
max :: TypeCheckError -> TypeCheckError -> TypeCheckError
$cmax :: TypeCheckError -> TypeCheckError -> TypeCheckError
>= :: TypeCheckError -> TypeCheckError -> Bool
$c>= :: TypeCheckError -> TypeCheckError -> Bool
> :: TypeCheckError -> TypeCheckError -> Bool
$c> :: TypeCheckError -> TypeCheckError -> Bool
<= :: TypeCheckError -> TypeCheckError -> Bool
$c<= :: TypeCheckError -> TypeCheckError -> Bool
< :: TypeCheckError -> TypeCheckError -> Bool
$c< :: TypeCheckError -> TypeCheckError -> Bool
compare :: TypeCheckError -> TypeCheckError -> Ordering
$ccompare :: TypeCheckError -> TypeCheckError -> Ordering
Ord, Int -> TypeCheckError -> ShowS
[TypeCheckError] -> ShowS
TypeCheckError -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TypeCheckError] -> ShowS
$cshowList :: [TypeCheckError] -> ShowS
show :: TypeCheckError -> String
$cshow :: TypeCheckError -> String
showsPrec :: Int -> TypeCheckError -> ShowS
$cshowsPrec :: Int -> TypeCheckError -> ShowS
Show)

-- State is closure environment, set of associations between type variables and types, unresolved type id supply
type AnnotateState = ExceptT TypeCheckError (State (PartialType, Set TypeAssociation, Int))

withNewEnv :: LocTag -> AnnotateState a -> AnnotateState (PartialType, a)
withNewEnv :: forall a.
LocTag -> AnnotateState a -> AnnotateState (PartialType, a)
withNewEnv LocTag
anno AnnotateState a
action = do
  (PartialType
env, Set TypeAssociation
typeMap, Int
v) <- forall s (m :: * -> *). MonadState s m => m s
State.get
  forall s (m :: * -> *). MonadState s m => s -> m ()
State.put (LocTag -> Int -> PartialType
TypeVariable LocTag
anno Int
v, Set TypeAssociation
typeMap, Int
v forall a. Num a => a -> a -> a
+ Int
1)
  a
result <- AnnotateState a
action
  forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
State.modify forall a b. (a -> b) -> a -> b
$ \(PartialType
_, Set TypeAssociation
typeMap, Int
v) -> (PartialType
env, Set TypeAssociation
typeMap, Int
v)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (LocTag -> Int -> PartialType
TypeVariable LocTag
anno Int
v, a
result)

setEnv :: PartialType -> AnnotateState ()
setEnv :: PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
setEnv PartialType
env = forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
State.modify forall a b. (a -> b) -> a -> b
$ \(PartialType
_, Set TypeAssociation
typeMap, Int
v) -> (PartialType
env, Set TypeAssociation
typeMap, Int
v)

data TypeAssociation = TypeAssociation Int PartialType
  deriving (TypeAssociation -> TypeAssociation -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TypeAssociation -> TypeAssociation -> Bool
$c/= :: TypeAssociation -> TypeAssociation -> Bool
== :: TypeAssociation -> TypeAssociation -> Bool
$c== :: TypeAssociation -> TypeAssociation -> Bool
Eq, Eq TypeAssociation
TypeAssociation -> TypeAssociation -> Bool
TypeAssociation -> TypeAssociation -> Ordering
TypeAssociation -> TypeAssociation -> TypeAssociation
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: TypeAssociation -> TypeAssociation -> TypeAssociation
$cmin :: TypeAssociation -> TypeAssociation -> TypeAssociation
max :: TypeAssociation -> TypeAssociation -> TypeAssociation
$cmax :: TypeAssociation -> TypeAssociation -> TypeAssociation
>= :: TypeAssociation -> TypeAssociation -> Bool
$c>= :: TypeAssociation -> TypeAssociation -> Bool
> :: TypeAssociation -> TypeAssociation -> Bool
$c> :: TypeAssociation -> TypeAssociation -> Bool
<= :: TypeAssociation -> TypeAssociation -> Bool
$c<= :: TypeAssociation -> TypeAssociation -> Bool
< :: TypeAssociation -> TypeAssociation -> Bool
$c< :: TypeAssociation -> TypeAssociation -> Bool
compare :: TypeAssociation -> TypeAssociation -> Ordering
$ccompare :: TypeAssociation -> TypeAssociation -> Ordering
Ord, Int -> TypeAssociation -> ShowS
[TypeAssociation] -> ShowS
TypeAssociation -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TypeAssociation] -> ShowS
$cshowList :: [TypeAssociation] -> ShowS
show :: TypeAssociation -> String
$cshow :: TypeAssociation -> String
showsPrec :: Int -> TypeAssociation -> ShowS
$cshowsPrec :: Int -> TypeAssociation -> ShowS
Show)

makeAssociations :: PartialType -> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations :: PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
ta PartialType
tb = case (PartialType
ta, PartialType
tb) of
  (PartialType
x, PartialType
y) | PartialType
x forall a. Eq a => a -> a -> Bool
== PartialType
y -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
  (PartialType
AnyType, PartialType
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
  (PartialType
_, PartialType
AnyType) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
  (TypeVariable LocTag
_ Int
i, PartialType
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Set a
Set.singleton forall a b. (a -> b) -> a -> b
$ Int -> PartialType -> TypeAssociation
TypeAssociation Int
i PartialType
tb
  (PartialType
_, TypeVariable LocTag
_ Int
i) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Set a
Set.singleton forall a b. (a -> b) -> a -> b
$ Int -> PartialType -> TypeAssociation
TypeAssociation Int
i PartialType
ta
  (ArrTypeP PartialType
a PartialType
b, ArrTypeP PartialType
c PartialType
d) -> forall a. Ord a => Set a -> Set a -> Set a
Set.union forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
a PartialType
c forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
b PartialType
d
  (PairTypeP PartialType
a PartialType
b, PairTypeP PartialType
c PartialType
d) -> forall a. Ord a => Set a -> Set a -> Set a
Set.union forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
a PartialType
c forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
b PartialType
d
  (PairTypeP PartialType
a PartialType
b, PartialType
ZeroTypeP) -> forall a. Ord a => Set a -> Set a -> Set a
Set.union forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
a PartialType
ZeroTypeP forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
b PartialType
ZeroTypeP
  (PartialType
ZeroTypeP, PairTypeP PartialType
a PartialType
b) -> forall a. Ord a => Set a -> Set a -> Set a
Set.union forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
a PartialType
ZeroTypeP forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
b PartialType
ZeroTypeP
  (PartialType, PartialType)
_ -> forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ PartialType -> PartialType -> TypeCheckError
InconsistentTypes PartialType
ta PartialType
tb

buildTypeMap :: Set TypeAssociation -> Either TypeCheckError (Map Int PartialType)
buildTypeMap :: Set TypeAssociation -> Either TypeCheckError (Map Int PartialType)
buildTypeMap Set TypeAssociation
assocSet =
  let multiMap :: Map Int (DList PartialType)
multiMap = forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith forall a. DList a -> DList a -> DList a
DList.append forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(TypeAssociation Int
i PartialType
t) -> (Int
i, forall a. a -> DList a
DList.singleton PartialType
t))
        forall a b. (a -> b) -> a -> b
$ forall a. Set a -> [a]
Set.toList Set TypeAssociation
assocSet
      getKeys :: PartialType -> DList Int
getKeys = \case
        TypeVariable LocTag
_ Int
i -> forall a. a -> DList a
DList.singleton Int
i
        ArrTypeP PartialType
a PartialType
b     -> PartialType -> DList Int
getKeys PartialType
a forall a. Semigroup a => a -> a -> a
<> PartialType -> DList Int
getKeys PartialType
b
        PairTypeP PartialType
a PartialType
b    -> PartialType -> DList Int
getKeys PartialType
a forall a. Semigroup a => a -> a -> a
<> PartialType -> DList Int
getKeys PartialType
b
        PartialType
_                -> forall a. Monoid a => a
mempty
      isRecursiveType :: Set Int -> Int -> Maybe Int
isRecursiveType Set Int
resolvedSet Int
k = case (forall a. Ord a => a -> Set a -> Bool
Set.member Int
k Set Int
resolvedSet, forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Int
k Map Int (DList PartialType)
multiMap) of
        (Bool
True, Maybe (DList PartialType)
_) -> forall a. a -> Maybe a
Just Int
k
        (Bool
_, Maybe (DList PartialType)
Nothing) -> forall a. Maybe a
Nothing
        (Bool
_, Just DList PartialType
t) -> forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Int
nk Maybe Int
r -> Set Int -> Int -> Maybe Int
isRecursiveType (forall a. Ord a => a -> Set a -> Set a
Set.insert Int
k Set Int
resolvedSet) Int
nk forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe Int
r) forall a. Maybe a
Nothing
          forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap PartialType -> DList Int
getKeys DList PartialType
t
      debugShowMap :: Map a a -> a -> a
debugShowMap Map a a
tm = forall a. String -> a -> a
debugTrace (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(a
k, a
v) -> forall a. Show a => a -> String
show a
k forall a. Semigroup a => a -> a -> a
<> String
": " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show a
v forall a. Semigroup a => a -> a -> a
<> String
"\n") forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
Map.toAscList Map a a
tm)
      buildMap :: Set TypeAssociation
-> Map Int PartialType
-> Either TypeCheckError (Map Int PartialType)
buildMap Set TypeAssociation
assoc Map Int PartialType
typeMap = case forall a. Set a -> Maybe (a, Set a)
Set.minView Set TypeAssociation
assoc of
        Maybe (TypeAssociation, Set TypeAssociation)
Nothing -> forall {a} {a} {a}. (Show a, Show a) => Map a a -> a -> a
debugShowMap Map Int PartialType
typeMap forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure Map Int PartialType
typeMap
        Just (TypeAssociation Int
i PartialType
t, Set TypeAssociation
newAssoc) -> case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Int
i Map Int PartialType
typeMap of
          Maybe PartialType
Nothing -> Set TypeAssociation
-> Map Int PartialType
-> Either TypeCheckError (Map Int PartialType)
buildMap Set TypeAssociation
newAssoc forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Int
i PartialType
t Map Int PartialType
typeMap
          Just PartialType
t2 -> PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
t PartialType
t2 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\Set TypeAssociation
assoc2 -> Set TypeAssociation
-> Map Int PartialType
-> Either TypeCheckError (Map Int PartialType)
buildMap (Set TypeAssociation
newAssoc forall a. Semigroup a => a -> a -> a
<> Set TypeAssociation
assoc2) Map Int PartialType
typeMap)
  -- if any variables result in lookup cycles, fail with RecursiveType
  in case forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Int
t Maybe Int
r -> Set Int -> Int -> Maybe Int
isRecursiveType forall a. Set a
Set.empty Int
t forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe Int
r) forall a. Maybe a
Nothing (forall k a. Map k a -> [k]
Map.keys Map Int (DList PartialType)
multiMap) of
    Just Int
k  -> forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ Int -> TypeCheckError
RecursiveType Int
k
    Maybe Int
Nothing -> forall a. String -> a -> a
debugTrace (forall a. Show a => a -> String
show Map Int (DList PartialType)
multiMap) forall a b. (a -> b) -> a -> b
$ Set TypeAssociation
-> Map Int PartialType
-> Either TypeCheckError (Map Int PartialType)
buildMap Set TypeAssociation
assocSet forall a. Monoid a => a
mempty

fullyResolve :: (Int -> Maybe PartialType) -> PartialType -> PartialType
fullyResolve :: (Int -> Maybe PartialType) -> PartialType -> PartialType
fullyResolve Int -> Maybe PartialType
resolve = PartialType -> PartialType
convert where
    convert :: PartialType -> PartialType
convert = forall a. Plated a => (a -> a) -> a -> a
transform PartialType -> PartialType
endo
    endo :: PartialType -> PartialType
endo = \case
      TypeVariable LocTag
anno Int
i -> case Int -> Maybe PartialType
resolve Int
i of
        Maybe PartialType
Nothing -> LocTag -> Int -> PartialType
TypeVariable LocTag
anno Int
i
        Just PartialType
t  -> PartialType -> PartialType
convert PartialType
t -- debugTrace (show t) $ convert t
      PartialType
x -> PartialType
x

traceAssociate :: PartialType -> PartialType -> a -> a
traceAssociate :: forall a. PartialType -> PartialType -> a -> a
traceAssociate PartialType
a PartialType
b = if Bool
debug
  then forall a. String -> a -> a
trace (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String
"associateVar ", forall a. Show a => a -> String
show PartialType
a, String
" -- ", forall a. Show a => a -> String
show PartialType
b])
  else forall a. a -> a
id

associateVar :: PartialType -> PartialType -> AnnotateState ()
associateVar :: PartialType
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateVar PartialType
a PartialType
b = forall e (m :: * -> *) a. MonadError e m => Either e a -> m a
liftEither (PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations PartialType
a PartialType
b) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Set TypeAssociation
set -> forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
State.modify (forall {b} {a} {c}. Semigroup b => b -> (a, b, c) -> (a, b, c)
changeState Set TypeAssociation
set) where
  changeState :: b -> (a, b, c) -> (a, b, c)
changeState b
set (a
curVar, b
oldSet, c
v) = (a
curVar, b
oldSet forall a. Semigroup a => a -> a -> a
<> b
set, c
v)

{-
-- convert a PartialType to a full type, aborting on circular references
fullyResolve :: Map Int PartialType -> PartialType -> Maybe DataType
fullyResolve typeMap x = case mostlyResolved of
  Left _ -> Nothing
  Right mr -> filterTypeVars mr
  where
    filterTypeVars ZeroTypeP = pure ZeroType
    filterTypeVars (TypeVariable _) = Nothing
    filterTypeVars (ArrTypeP a b) = ArrType <$> filterTypeVars a <*> filterTypeVars b
    filterTypeVars (PairTypeP a b) = PairType <$> filterTypeVars a <*> filterTypeVars b
    mostlyResolved = mostlyResolve typeMap x
-}

{--
 - reserve 0 -> n*2 TypeVariables for types of FragExprs
 -}
initState :: Term3 -> (PartialType, Set TypeAssociation, Int)
initState :: Term3 -> (PartialType, Set TypeAssociation, Int)
initState (Term3 Map FragIndex FragExprUR
termMap) =
  let startVariable :: Int
startVariable = case forall a. Set a -> Maybe (a, Set a)
Set.maxView (forall k a. Map k a -> Set k
Map.keysSet Map FragIndex FragExprUR
termMap) of
        Maybe (FragIndex, Set FragIndex)
Nothing               -> Int
0
        Just (FragIndex Int
i, Set FragIndex
_) -> (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a. Num a => a -> a -> a
* Int
2
  in (LocTag -> Int -> PartialType
TypeVariable LocTag
DummyLoc Int
0, forall a. Set a
Set.empty, Int
startVariable)

getFragType :: LocTag -> FragIndex -> PartialType
getFragType :: LocTag -> FragIndex -> PartialType
getFragType LocTag
anno (FragIndex Int
i) = PartialType -> PartialType -> PartialType
ArrTypeP (LocTag -> Int -> PartialType
TypeVariable LocTag
anno forall a b. (a -> b) -> a -> b
$ Int
i forall a. Num a => a -> a -> a
* Int
2) (LocTag -> Int -> PartialType
TypeVariable LocTag
anno forall a b. (a -> b) -> a -> b
$ Int
i forall a. Num a => a -> a -> a
* Int
2 forall a. Num a => a -> a -> a
+ Int
1)

annotate :: Term3 -> AnnotateState PartialType
annotate :: Term3 -> AnnotateState PartialType
annotate (Term3 Map FragIndex FragExprUR
termMap) =
  let annotate' :: Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag -> AnnotateState PartialType
      annotate' :: Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' = \case
        LocTag
anno :< FragExprF
  (RecursionSimulationPieces FragExprUR)
  (Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag)
ZeroFragF -> forall (f :: * -> *) a. Applicative f => a -> f a
pure PartialType
ZeroTypeP
        LocTag
anno :< PairFragF Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
a Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
b -> PartialType -> PartialType -> PartialType
PairTypeP forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
a forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
b
        LocTag
anno :< FragExprF
  (RecursionSimulationPieces FragExprUR)
  (Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag)
EnvFragF -> (forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
State.gets (\(PartialType
t, Set TypeAssociation
_, Int
_) -> PartialType
t))
        LocTag
anno :< SetEnvFragF Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x -> do
          PartialType
xt <- Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x
          (PartialType
it, (PartialType
ot, ()
_)) <- forall a.
LocTag -> AnnotateState a -> AnnotateState (PartialType, a)
withNewEnv LocTag
anno forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a.
LocTag -> AnnotateState a -> AnnotateState (PartialType, a)
withNewEnv LocTag
anno forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          PartialType
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateVar (PartialType -> PartialType -> PartialType
PairTypeP (PartialType -> PartialType -> PartialType
ArrTypeP PartialType
it PartialType
ot) PartialType
it) PartialType
xt
          forall (f :: * -> *) a. Applicative f => a -> f a
pure PartialType
ot
        LocTag
anno :< DeferFragF FragIndex
ind -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ LocTag -> FragIndex -> PartialType
getFragType LocTag
anno FragIndex
ind
        LocTag
anno :< FragExprF
  (RecursionSimulationPieces FragExprUR)
  (Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag)
AbortFragF -> do
          (PartialType
it, ()
_) <- forall a.
LocTag -> AnnotateState a -> AnnotateState (PartialType, a)
withNewEnv LocTag
anno forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (PartialType -> PartialType -> PartialType
ArrTypeP PartialType
ZeroTypeP (PartialType -> PartialType -> PartialType
ArrTypeP PartialType
it PartialType
it))
        LocTag
anno :< GateFragF Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
l Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
r -> do
          PartialType
lt <- Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
l
          PartialType
rt <- Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
r
          PartialType
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateVar PartialType
lt PartialType
rt
          forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ PartialType -> PartialType -> PartialType
ArrTypeP PartialType
ZeroTypeP PartialType
lt
        LocTag
anno :< LeftFragF Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x -> do
          PartialType
xt <- Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x
          (PartialType
la, ()
_) <- forall a.
LocTag -> AnnotateState a -> AnnotateState (PartialType, a)
withNewEnv LocTag
anno forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          PartialType
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateVar (PartialType -> PartialType -> PartialType
PairTypeP PartialType
la PartialType
AnyType) PartialType
xt
          forall (f :: * -> *) a. Applicative f => a -> f a
pure PartialType
la
        LocTag
anno :< RightFragF Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x -> do
          PartialType
xt <- Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x
          (PartialType
ra, ()
_) <- forall a.
LocTag -> AnnotateState a -> AnnotateState (PartialType, a)
withNewEnv LocTag
anno forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          PartialType
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateVar (PartialType -> PartialType -> PartialType
PairTypeP PartialType
AnyType PartialType
ra) PartialType
xt
          forall (f :: * -> *) a. Applicative f => a -> f a
pure PartialType
ra
        LocTag
anno :< FragExprF
  (RecursionSimulationPieces FragExprUR)
  (Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag)
TraceFragF -> (forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
State.gets (\(PartialType
t, Set TypeAssociation
_, Int
_) -> PartialType
t))
        LocTag
anno :< AuxFragF (NestedSetEnvs UnsizedRecursionToken
_) -> (forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
State.gets (\(PartialType
t, Set TypeAssociation
_, Int
_) -> PartialType
t))
        LocTag
anno :< AuxFragF (SizingWrapper UnsizedRecursionToken
_ (FragExprUR Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x)) -> Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
x
      initInputType :: LocTag -> FragIndex -> AnnotateState ()
      initInputType :: LocTag
-> FragIndex
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
initInputType LocTag
anno FragIndex
fi = let (ArrTypeP PartialType
it PartialType
_) = LocTag -> FragIndex -> PartialType
getFragType LocTag
anno FragIndex
fi in forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
State.modify (\(PartialType
_, Set TypeAssociation
s, Int
i) -> (PartialType
it, Set TypeAssociation
s, Int
i))
      associateOutType :: LocTag -> FragIndex -> PartialType -> AnnotateState ()
      associateOutType :: LocTag
-> FragIndex
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateOutType LocTag
anno FragIndex
fi PartialType
ot = let (ArrTypeP PartialType
_ PartialType
ot2) = LocTag -> FragIndex -> PartialType
getFragType LocTag
anno FragIndex
fi in PartialType
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateVar PartialType
ot PartialType
ot2
      rootType :: LocTag -> AnnotateState PartialType
rootType LocTag
anno = LocTag
-> FragIndex
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
initInputType LocTag
anno (Int -> FragIndex
FragIndex Int
0) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' (FragExprUR
-> Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
unFragExprUR forall a b. (a -> b) -> a -> b
$ forall a. Map FragIndex a -> a
rootFrag Map FragIndex FragExprUR
termMap)
  in forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ (forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (\FragIndex
k FragExprUR
v -> LocTag
-> FragIndex
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
initInputType LocTag
DummyLoc FragIndex
k forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
-> AnnotateState PartialType
annotate' (FragExprUR
-> Cofree (FragExprF (RecursionSimulationPieces FragExprUR)) LocTag
unFragExprUR FragExprUR
v) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= LocTag
-> FragIndex
-> PartialType
-> ExceptT
     TypeCheckError (State (PartialType, Set TypeAssociation, Int)) ()
associateOutType LocTag
DummyLoc FragIndex
k) Map FragIndex FragExprUR
termMap) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> LocTag -> AnnotateState PartialType
rootType LocTag
DummyLoc

partiallyAnnotate :: Term3 -> Either TypeCheckError (PartialType, Int -> Maybe PartialType)
partiallyAnnotate :: Term3
-> Either TypeCheckError (PartialType, Int -> Maybe PartialType)
partiallyAnnotate Term3
term =
  let runner :: State (PartialType, Set TypeAssociation, Int) (Either TypeCheckError PartialType)
      runner :: State
  (PartialType, Set TypeAssociation, Int)
  (Either TypeCheckError PartialType)
runner = forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT forall a b. (a -> b) -> a -> b
$ Term3 -> AnnotateState PartialType
annotate Term3
term
      (Either TypeCheckError PartialType
rt, (PartialType
_, Set TypeAssociation
s, Int
_)) = forall s a. State s a -> s -> (a, s)
State.runState State
  (PartialType, Set TypeAssociation, Int)
  (Either TypeCheckError PartialType)
runner (Term3 -> (PartialType, Set TypeAssociation, Int)
initState Term3
term)
  in (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either TypeCheckError PartialType
rt forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set TypeAssociation -> Either TypeCheckError (Map Int PartialType)
buildTypeMap Set TypeAssociation
s)

inferType :: Term3 -> Either TypeCheckError PartialType
inferType :: Term3 -> Either TypeCheckError PartialType
inferType Term3
tm = (PartialType, Int -> Maybe PartialType) -> PartialType
lookupFully forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term3
-> Either TypeCheckError (PartialType, Int -> Maybe PartialType)
partiallyAnnotate Term3
tm where
  lookupFully :: (PartialType, Int -> Maybe PartialType) -> PartialType
lookupFully (PartialType
ty, Int -> Maybe PartialType
resolver) = (Int -> Maybe PartialType) -> PartialType -> PartialType
fullyResolve Int -> Maybe PartialType
resolver PartialType
ty

typeCheck :: PartialType -> Term3 -> Maybe TypeCheckError
typeCheck :: PartialType -> Term3 -> Maybe TypeCheckError
typeCheck PartialType
t tm :: Term3
tm@(Term3 Map FragIndex FragExprUR
typeMap) = forall {a} {b}. Either a b -> Maybe a
convert (Term3
-> Either TypeCheckError (PartialType, Int -> Maybe PartialType)
partiallyAnnotate Term3
tm forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (PartialType, Int -> Maybe PartialType)
-> Either TypeCheckError (Set TypeAssociation)
associate) where
  associate :: (PartialType, Int -> Maybe PartialType)
-> Either TypeCheckError (Set TypeAssociation)
associate (PartialType
ty, Int -> Maybe PartialType
resolver) = forall a. String -> a -> a
debugTrace (String
"COMPARING TYPES " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (PartialType
t, (Int -> Maybe PartialType) -> PartialType -> PartialType
fullyResolve Int -> Maybe PartialType
resolver PartialType
ty) forall a. Semigroup a => a -> a -> a
<> String
"\n" forall a. Semigroup a => a -> a -> a
<> PartialType -> (Int -> Maybe PartialType) -> String
debugMap PartialType
ty Int -> Maybe PartialType
resolver)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a}. Show a => a -> a
traceAgain forall a b. (a -> b) -> a -> b
$ PartialType
-> PartialType -> Either TypeCheckError (Set TypeAssociation)
makeAssociations ((Int -> Maybe PartialType) -> PartialType -> PartialType
fullyResolve Int -> Maybe PartialType
resolver PartialType
ty) PartialType
t
  debugMap :: PartialType -> (Int -> Maybe PartialType) -> String
debugMap PartialType
ty Int -> Maybe PartialType
resolver = TypeDebugInfo -> String
showTypeDebugInfo (Term3 -> (FragIndex -> PartialType) -> PartialType -> TypeDebugInfo
TypeDebugInfo Term3
tm ((Int -> Maybe PartialType) -> PartialType -> PartialType
fullyResolve Int -> Maybe PartialType
resolver forall b c a. (b -> c) -> (a -> b) -> a -> c
. LocTag -> FragIndex -> PartialType
getFragType LocTag
DummyLoc)
                                            ((Int -> Maybe PartialType) -> PartialType -> PartialType
fullyResolve Int -> Maybe PartialType
resolver PartialType
ty))
  traceAgain :: a -> a
traceAgain a
s = forall a. String -> a -> a
debugTrace (String
"Resulting thing " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show a
s) a
s
  convert :: Either a b -> Maybe a
convert = \case
    Left a
er -> forall a. a -> Maybe a
Just a
er
    Either a b
_       -> forall a. Maybe a
Nothing