{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}

module U.Util.Type where

import Data.Set (Set)
import qualified Data.Set as Set
import U.Codebase.Type (F' (..), TypeR)
import U.Core.ABT (pattern Var')
import qualified U.Core.ABT as ABT
import qualified U.Core.ABT.Var as ABT

-- * Constructors

effect :: (Ord v) => [TypeR r v] -> TypeR r v -> TypeR r v
effect :: forall v r. Ord v => [TypeR r v] -> TypeR r v -> TypeR r v
effect [TypeR r v]
es (Effect1' TypeR r v
fs TypeR r v
t) =
  let es' :: [TypeR r v]
es' = ([TypeR r v]
es [TypeR r v] -> (TypeR r v -> [TypeR r v]) -> [TypeR r v]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= TypeR r v -> [TypeR r v]
forall r v. TypeR r v -> [TypeR r v]
flattenEffects) [TypeR r v] -> [TypeR r v] -> [TypeR r v]
forall a. [a] -> [a] -> [a]
++ TypeR r v -> [TypeR r v]
forall r v. TypeR r v -> [TypeR r v]
flattenEffects TypeR r v
fs
   in () -> F' r (TypeR r v) -> TypeR r v
forall (f :: * -> *) v a.
(Foldable f, Ord v) =>
a -> f (Term f v a) -> Term f v a
ABT.tm () (TypeR r v -> TypeR r v -> F' r (TypeR r v)
forall r a. a -> a -> F' r a
Effect (() -> F' r (TypeR r v) -> TypeR r v
forall (f :: * -> *) v a.
(Foldable f, Ord v) =>
a -> f (Term f v a) -> Term f v a
ABT.tm () ([TypeR r v] -> F' r (TypeR r v)
forall r a. [a] -> F' r a
Effects [TypeR r v]
es')) TypeR r v
t)
effect [TypeR r v]
es TypeR r v
t = () -> F' r (TypeR r v) -> TypeR r v
forall (f :: * -> *) v a.
(Foldable f, Ord v) =>
a -> f (Term f v a) -> Term f v a
ABT.tm () (TypeR r v -> TypeR r v -> F' r (TypeR r v)
forall r a. a -> a -> F' r a
Effect (() -> F' r (TypeR r v) -> TypeR r v
forall (f :: * -> *) v a.
(Foldable f, Ord v) =>
a -> f (Term f v a) -> Term f v a
ABT.tm () ([TypeR r v] -> F' r (TypeR r v)
forall r a. [a] -> F' r a
Effects [TypeR r v]
es)) TypeR r v
t)

effects :: (Ord v) => [TypeR r v] -> TypeR r v
effects :: forall v r. Ord v => [TypeR r v] -> TypeR r v
effects [TypeR r v]
es = () -> F' r (TypeR r v) -> TypeR r v
forall (f :: * -> *) v a.
(Foldable f, Ord v) =>
a -> f (Term f v a) -> Term f v a
ABT.tm () ([TypeR r v] -> F' r (TypeR r v)
forall r a. [a] -> F' r a
Effects ([TypeR r v] -> F' r (TypeR r v))
-> [TypeR r v] -> F' r (TypeR r v)
forall a b. (a -> b) -> a -> b
$ [TypeR r v]
es [TypeR r v] -> (TypeR r v -> [TypeR r v]) -> [TypeR r v]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= TypeR r v -> [TypeR r v]
forall r v. TypeR r v -> [TypeR r v]
flattenEffects)

-- * Modification

-- Remove all effect variables from the type.
-- Used for type-based search, we apply this transformation to both the
-- indexed type and the query type, so the user can supply `a -> b` that will
-- match `a ->{e} b` (but not `a ->{IO} b`).
removeAllEffectVars :: (ABT.Var v) => TypeR r v -> TypeR r v
removeAllEffectVars :: forall v r. Var v => TypeR r v -> TypeR r v
removeAllEffectVars TypeR r v
t =
  let allEffectVars :: Set v
allEffectVars = (TypeR r v -> Set v) -> [TypeR r v] -> Set v
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap TypeR r v -> Set v
forall {v} {r}. Ord v => TypeR r v -> Set v
go (TypeR r v -> [TypeR r v]
forall v (f :: * -> *) a.
(Ord v, Traversable f) =>
Term f v a -> [Term f v a]
ABT.subterms TypeR r v
t)
      go :: TypeR r v -> Set v
go (Effects' [TypeR r v]
vs) = [v] -> Set v
forall a. Ord a => [a] -> Set a
Set.fromList [v
v | Var' v
v <- [TypeR r v]
vs]
      go (Effect1' (Var' v
v) TypeR r v
_) = v -> Set v
forall a. a -> Set a
Set.singleton v
v
      go TypeR r v
_ = Set v
forall a. Monoid a => a
mempty
      ([v]
vs, TypeR r v
tu) = TypeR r v -> ([v], TypeR r v)
forall r v. TypeR r v -> ([v], TypeR r v)
unforall' TypeR r v
t
   in [v] -> TypeR r v -> TypeR r v
forall v r. Ord v => [v] -> TypeR r v -> TypeR r v
generalize [v]
vs (Set v -> TypeR r v -> TypeR r v
forall v r. Var v => Set v -> TypeR r v -> TypeR r v
removeEffectVars Set v
allEffectVars TypeR r v
tu)

-- Remove free effect variables from the type that are in the set
removeEffectVars :: (ABT.Var v) => Set v -> TypeR r v -> TypeR r v
removeEffectVars :: forall v r. Var v => Set v -> TypeR r v -> TypeR r v
removeEffectVars Set v
removals TypeR r v
t =
  let z :: TypeR r v
z = [TypeR r v] -> TypeR r v
forall v r. Ord v => [TypeR r v] -> TypeR r v
effects []
      t' :: TypeR r v
t' = [(v, TypeR r v)] -> TypeR r v -> TypeR r v
forall (f :: * -> *) v b a.
(Foldable f, Functor f, Var v) =>
[(v, Term f v b)] -> Term f v a -> Term f v a
ABT.substsInheritAnnotation ((,TypeR r v
forall {r}. TypeR r v
z) (v -> (v, TypeR r v)) -> [v] -> [(v, TypeR r v)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Set v -> [v]
forall a. Set a -> [a]
Set.toList Set v
removals) TypeR r v
t
      -- leave explicitly empty `{}` alone
      removeEmpty :: Term (F' r) v () -> Maybe (Term (F' r) v ())
removeEmpty (Effect1' (Effects' []) Term (F' r) v ()
v) = Term (F' r) v () -> Maybe (Term (F' r) v ())
forall a. a -> Maybe a
Just ((Term (F' r) v () -> Maybe (Term (F' r) v ()))
-> Term (F' r) v () -> Term (F' r) v ()
forall (f :: * -> *) v a.
(Traversable f, Ord v) =>
(Term f v a -> Maybe (Term f v a)) -> Term f v a -> Term f v a
ABT.visitPure Term (F' r) v () -> Maybe (Term (F' r) v ())
removeEmpty Term (F' r) v ()
v)
      removeEmpty (Effect1' Term (F' r) v ()
e Term (F' r) v ()
v) =
        case Term (F' r) v () -> [Term (F' r) v ()]
forall r v. TypeR r v -> [TypeR r v]
flattenEffects Term (F' r) v ()
e of
          [] -> Term (F' r) v () -> Maybe (Term (F' r) v ())
forall a. a -> Maybe a
Just ((Term (F' r) v () -> Maybe (Term (F' r) v ()))
-> Term (F' r) v () -> Term (F' r) v ()
forall (f :: * -> *) v a.
(Traversable f, Ord v) =>
(Term f v a -> Maybe (Term f v a)) -> Term f v a -> Term f v a
ABT.visitPure Term (F' r) v () -> Maybe (Term (F' r) v ())
removeEmpty Term (F' r) v ()
v)
          [Term (F' r) v ()]
es -> Term (F' r) v () -> Maybe (Term (F' r) v ())
forall a. a -> Maybe a
Just ([Term (F' r) v ()] -> Term (F' r) v () -> Term (F' r) v ()
forall v r. Ord v => [TypeR r v] -> TypeR r v -> TypeR r v
effect [Term (F' r) v ()]
es (Term (F' r) v () -> Term (F' r) v ())
-> Term (F' r) v () -> Term (F' r) v ()
forall a b. (a -> b) -> a -> b
$ (Term (F' r) v () -> Maybe (Term (F' r) v ()))
-> Term (F' r) v () -> Term (F' r) v ()
forall (f :: * -> *) v a.
(Traversable f, Ord v) =>
(Term f v a -> Maybe (Term f v a)) -> Term f v a -> Term f v a
ABT.visitPure Term (F' r) v () -> Maybe (Term (F' r) v ())
removeEmpty Term (F' r) v ()
v)
      removeEmpty (Effects' [Term (F' r) v ()]
es) =
        Term (F' r) v () -> Maybe (Term (F' r) v ())
forall a. a -> Maybe a
Just (Term (F' r) v () -> Maybe (Term (F' r) v ()))
-> Term (F' r) v () -> Maybe (Term (F' r) v ())
forall a b. (a -> b) -> a -> b
$ [Term (F' r) v ()] -> Term (F' r) v ()
forall v r. Ord v => [TypeR r v] -> TypeR r v
effects ([Term (F' r) v ()]
es [Term (F' r) v ()]
-> (Term (F' r) v () -> [Term (F' r) v ()]) -> [Term (F' r) v ()]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Term (F' r) v () -> [Term (F' r) v ()]
forall r v. TypeR r v -> [TypeR r v]
flattenEffects)
      removeEmpty Term (F' r) v ()
_ = Maybe (Term (F' r) v ())
forall a. Maybe a
Nothing
   in (TypeR r v -> Maybe (TypeR r v)) -> TypeR r v -> TypeR r v
forall (f :: * -> *) v a.
(Traversable f, Ord v) =>
(Term f v a -> Maybe (Term f v a)) -> Term f v a -> Term f v a
ABT.visitPure TypeR r v -> Maybe (TypeR r v)
forall {v} {r}.
Ord v =>
Term (F' r) v () -> Maybe (Term (F' r) v ())
removeEmpty TypeR r v
t'

flattenEffects :: TypeR r v -> [TypeR r v]
flattenEffects :: forall r v. TypeR r v -> [TypeR r v]
flattenEffects (Effects' [TypeR r v]
es) = [TypeR r v]
es [TypeR r v] -> (TypeR r v -> [TypeR r v]) -> [TypeR r v]
forall a b. [a] -> (a -> [b]) -> [b]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= TypeR r v -> [TypeR r v]
forall r v. TypeR r v -> [TypeR r v]
flattenEffects
flattenEffects TypeR r v
es = [TypeR r v
es]

-- | Bind the given variables with an outer `forall`, if they are used in `t`.
generalize :: (Ord v) => [v] -> TypeR r v -> TypeR r v
generalize :: forall v r. Ord v => [v] -> TypeR r v -> TypeR r v
generalize [v]
vs TypeR r v
t = (v -> TypeR r v -> TypeR r v) -> TypeR r v -> [v] -> TypeR r v
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr v -> TypeR r v -> TypeR r v
forall {v} {r}. Ord v => v -> Term (F' r) v () -> Term (F' r) v ()
f TypeR r v
t [v]
vs
  where
    f :: v -> Term (F' r) v () -> Term (F' r) v ()
f v
v Term (F' r) v ()
t = if v -> Set v -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member v
v (Term (F' r) v () -> Set v
forall (f :: * -> *) v a. Term f v a -> Set v
ABT.freeVars Term (F' r) v ()
t) then v -> Term (F' r) v () -> Term (F' r) v ()
forall {v} {r}. Ord v => v -> Term (F' r) v () -> Term (F' r) v ()
forAll v
v Term (F' r) v ()
t else Term (F' r) v ()
t

-- * Patterns

pattern ForallsNamed' :: [v] -> TypeR r v -> TypeR r v
pattern $mForallsNamed' :: forall {r} {v} {r}.
TypeR r v -> ([v] -> TypeR r v -> r) -> ((# #) -> r) -> r
ForallsNamed' vs body <- (unForalls -> Just (vs, body))

pattern ForallNamed' :: v -> TypeR r v -> TypeR r v
pattern $mForallNamed' :: forall {r} {v} {r}.
TypeR r v -> (v -> TypeR r v -> r) -> ((# #) -> r) -> r
ForallNamed' v body <- ABT.Tm' (Forall (ABT.out -> ABT.Abs v body))

pattern Effects' :: [TypeR r v] -> TypeR r v
pattern $mEffects' :: forall {r} {r} {v}.
TypeR r v -> ([TypeR r v] -> r) -> ((# #) -> r) -> r
Effects' es <- ABT.Tm' (Effects es)

pattern Effect1' :: TypeR r v -> TypeR r v -> TypeR r v
pattern $mEffect1' :: forall {r} {r} {v}.
TypeR r v -> (TypeR r v -> TypeR r v -> r) -> ((# #) -> r) -> r
Effect1' e t <- ABT.Tm' (Effect e t)

pattern Ref' :: r -> TypeR r v
pattern $mRef' :: forall {r} {r} {v}. TypeR r v -> (r -> r) -> ((# #) -> r) -> r
Ref' r <- ABT.Tm' (Ref r)

forAll :: (Ord v) => v -> TypeR r v -> TypeR r v
forAll :: forall {v} {r}. Ord v => v -> Term (F' r) v () -> Term (F' r) v ()
forAll v
v TypeR r v
body = () -> F' r (TypeR r v) -> TypeR r v
forall (f :: * -> *) v a.
(Foldable f, Ord v) =>
a -> f (Term f v a) -> Term f v a
ABT.tm () (TypeR r v -> F' r (TypeR r v)
forall r a. a -> F' r a
Forall (() -> v -> TypeR r v -> TypeR r v
forall v a (f :: * -> *).
Ord v =>
a -> v -> Term f v a -> Term f v a
ABT.abs () v
v TypeR r v
body))

unforall' :: TypeR r v -> ([v], TypeR r v)
unforall' :: forall r v. TypeR r v -> ([v], TypeR r v)
unforall' (ForallsNamed' [v]
vs TypeR r v
t) = ([v]
vs, TypeR r v
t)
unforall' TypeR r v
t = ([], TypeR r v
t)

unForalls :: TypeR r v -> Maybe ([v], TypeR r v)
unForalls :: forall r v. TypeR r v -> Maybe ([v], TypeR r v)
unForalls TypeR r v
t = TypeR r v -> [v] -> Maybe ([v], TypeR r v)
forall {r} {a}. TypeR r a -> [a] -> Maybe ([a], TypeR r a)
go TypeR r v
t []
  where
    go :: TypeR r a -> [a] -> Maybe ([a], TypeR r a)
go (ForallNamed' a
v TypeR r a
body) [a]
vs = TypeR r a -> [a] -> Maybe ([a], TypeR r a)
go TypeR r a
body (a
v a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
vs)
    go TypeR r a
_body [] = Maybe ([a], TypeR r a)
forall a. Maybe a
Nothing
    go TypeR r a
body [a]
vs = ([a], TypeR r a) -> Maybe ([a], TypeR r a)
forall a. a -> Maybe a
Just ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
vs, TypeR r a
body)