-- | Functions for working with term and type annotations

module Hydra.Annotations where

import Hydra.Basics
import Hydra.Strip
import Hydra.Core
import Hydra.Compute
import Hydra.Extras
import Hydra.Graph
import Hydra.CoreDecoding
import Hydra.CoreEncoding
import Hydra.Mantle
import Hydra.Rewriting
import Hydra.Tier1
import Hydra.Tier2
import qualified Hydra.Dsl.Expect as Expect
import qualified Hydra.Dsl.Terms as Terms

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


key_classes :: Name
key_classes = String -> Name
Name String
"classes"
key_description :: Name
key_description = String -> Name
Name String
"description"
key_type :: Name
key_type = String -> Name
Name String
"type"


aggregateAnnotations :: (x -> Maybe y) -> (y -> x) -> (y -> M.Map Name Term) -> x -> M.Map Name Term
aggregateAnnotations :: forall x y.
(x -> Maybe y)
-> (y -> x) -> (y -> Map Name Term) -> x -> Map Name Term
aggregateAnnotations x -> Maybe y
getValue y -> x
getX y -> Map Name Term
getAnns x
t = [(Name, Term)] -> Map Name Term
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, Term)] -> Map Name Term)
-> [(Name, Term)] -> Map Name Term
forall a b. (a -> b) -> a -> b
$ [[(Name, Term)]] -> [(Name, Term)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
L.concat ([[(Name, Term)]] -> [(Name, Term)])
-> [[(Name, Term)]] -> [(Name, Term)]
forall a b. (a -> b) -> a -> b
$ [[(Name, Term)]] -> x -> [[(Name, Term)]]
toPairs [] x
t
  where
    toPairs :: [[(Name, Term)]] -> x -> [[(Name, Term)]]
toPairs [[(Name, Term)]]
rest x
t = case x -> Maybe y
getValue x
t of
      Maybe y
Nothing -> [[(Name, Term)]]
rest
      Just y
yy -> [[(Name, Term)]] -> x -> [[(Name, Term)]]
toPairs ((Map Name Term -> [(Name, Term)]
forall k a. Map k a -> [(k, a)]
M.toList (y -> Map Name Term
getAnns y
yy))[(Name, Term)] -> [[(Name, Term)]] -> [[(Name, Term)]]
forall a. a -> [a] -> [a]
:[[(Name, Term)]]
rest) (y -> x
getX y
yy)

failOnFlag :: Name -> String -> Flow s ()
failOnFlag :: forall s. Name -> String -> Flow s ()
failOnFlag Name
flag String
msg = do
  Bool
val <- Name -> Flow s Bool
forall s. Name -> Flow s Bool
hasFlag Name
flag
  if Bool
val
    then String -> Flow s ()
forall a. String -> Flow s a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
msg
    else () -> Flow s ()
forall a. a -> Flow s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

getAttr :: Name -> Flow s (Maybe Term)
getAttr :: forall s. Name -> Flow s (Maybe Term)
getAttr Name
key = (s -> Trace -> FlowState s (Maybe Term)) -> Flow s (Maybe Term)
forall s x. (s -> Trace -> FlowState s x) -> Flow s x
Flow s -> Trace -> FlowState s (Maybe Term)
forall {s}. s -> Trace -> FlowState s (Maybe Term)
q
  where
    q :: s -> Trace -> FlowState s (Maybe Term)
q s
s0 Trace
t0 = Maybe (Maybe Term) -> s -> Trace -> FlowState s (Maybe Term)
forall s x. Maybe x -> s -> Trace -> FlowState s x
FlowState (Maybe Term -> Maybe (Maybe Term)
forall a. a -> Maybe a
Just (Maybe Term -> Maybe (Maybe Term))
-> Maybe Term -> Maybe (Maybe Term)
forall a b. (a -> b) -> a -> b
$ Name -> Map Name Term -> Maybe Term
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
key (Map Name Term -> Maybe Term) -> Map Name Term -> Maybe Term
forall a b. (a -> b) -> a -> b
$ Trace -> Map Name Term
traceOther Trace
t0) s
s0 Trace
t0

getAttrWithDefault :: Name -> Term -> Flow s Term
getAttrWithDefault :: forall s. Name -> Term -> Flow s Term
getAttrWithDefault Name
key Term
def = Term -> Maybe Term -> Term
forall a. a -> Maybe a -> a
Y.fromMaybe Term
def (Maybe Term -> Term) -> Flow s (Maybe Term) -> Flow s Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Flow s (Maybe Term)
forall s. Name -> Flow s (Maybe Term)
getAttr Name
key

getDescription :: M.Map Name Term -> Flow Graph (Y.Maybe String)
getDescription :: Map Name Term -> Flow Graph (Maybe String)
getDescription Map Name Term
anns = case Name -> Map Name Term -> Maybe Term
getAnnotation Name
key_description Map Name Term
anns of
  Maybe Term
Nothing -> Maybe String -> Flow Graph (Maybe String)
forall a. a -> Flow Graph a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe String
forall a. Maybe a
Nothing
  Just Term
term -> case Term
term of
    TermLiteral (LiteralString String
s) -> Maybe String -> Flow Graph (Maybe String)
forall a. a -> Flow Graph a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe String -> Flow Graph (Maybe String))
-> Maybe String -> Flow Graph (Maybe String)
forall a b. (a -> b) -> a -> b
$ String -> Maybe String
forall a. a -> Maybe a
Just String
s
    Term
_ -> String -> Flow Graph (Maybe String)
forall a. String -> Flow Graph a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Flow Graph (Maybe String))
-> String -> Flow Graph (Maybe String)
forall a b. (a -> b) -> a -> b
$ String
"unexpected value for " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
key_description String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Term -> String
forall a. Show a => a -> String
show Term
term

getTermAnnotation :: Name -> Term -> Y.Maybe Term
getTermAnnotation :: Name -> Term -> Maybe Term
getTermAnnotation Name
key = Name -> Map Name Term -> Maybe Term
getAnnotation Name
key (Map Name Term -> Maybe Term)
-> (Term -> Map Name Term) -> Term -> Maybe Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Map Name Term
termAnnotationInternal

getTermDescription :: Term -> Flow Graph (Y.Maybe String)
getTermDescription :: Term -> Flow Graph (Maybe String)
getTermDescription = Map Name Term -> Flow Graph (Maybe String)
getDescription (Map Name Term -> Flow Graph (Maybe String))
-> (Term -> Map Name Term) -> Term -> Flow Graph (Maybe String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Map Name Term
termAnnotationInternal

getType :: M.Map Name Term -> Flow Graph (Y.Maybe Type)
getType :: Map Name Term -> Flow Graph (Maybe Type)
getType Map Name Term
anns = case Name -> Map Name Term -> Maybe Term
getAnnotation Name
key_type Map Name Term
anns of
  Maybe Term
Nothing -> Maybe Type -> Flow Graph (Maybe Type)
forall a. a -> Flow Graph a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Type
forall a. Maybe a
Nothing
  Just Term
dat -> Type -> Maybe Type
forall a. a -> Maybe a
Just (Type -> Maybe Type) -> Flow Graph Type -> Flow Graph (Maybe Type)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term -> Flow Graph Type
coreDecodeType Term
dat

getTypeAnnotation :: Name -> Type -> Y.Maybe Term
getTypeAnnotation :: Name -> Type -> Maybe Term
getTypeAnnotation Name
key = Name -> Map Name Term -> Maybe Term
getAnnotation Name
key (Map Name Term -> Maybe Term)
-> (Type -> Map Name Term) -> Type -> Maybe Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Map Name Term
typeAnnotationInternal

getTypeClasses :: Type -> Flow Graph (M.Map Name (S.Set TypeClass))
getTypeClasses :: Type -> Flow Graph (Map Name (Set TypeClass))
getTypeClasses Type
typ = case Name -> Type -> Maybe Term
getTypeAnnotation Name
key_classes Type
typ of
    Maybe Term
Nothing -> Map Name (Set TypeClass) -> Flow Graph (Map Name (Set TypeClass))
forall a. a -> Flow Graph a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map Name (Set TypeClass)
forall k a. Map k a
M.empty
    Just Term
term -> (Term -> Flow Graph Name)
-> (Term -> Flow Graph (Set TypeClass))
-> Term
-> Flow Graph (Map Name (Set TypeClass))
forall k s v.
Ord k =>
(Term -> Flow s k)
-> (Term -> Flow s v) -> Term -> Flow s (Map k v)
Expect.map Term -> Flow Graph Name
coreDecodeName ((Term -> Flow Graph TypeClass)
-> Term -> Flow Graph (Set TypeClass)
forall x s. Ord x => (Term -> Flow s x) -> Term -> Flow s (Set x)
Expect.set Term -> Flow Graph TypeClass
forall {s}. Term -> Flow s TypeClass
decodeClass) Term
term
  where
    decodeClass :: Term -> Flow s TypeClass
decodeClass Term
term = Name -> Term -> Flow s Name
forall s. Name -> Term -> Flow s Name
Expect.unitVariant Name
_TypeClass Term
term Flow s Name -> (Name -> Flow s TypeClass) -> Flow s TypeClass
forall a b. Flow s a -> (a -> Flow s b) -> Flow s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Name
fn -> Flow s TypeClass
-> (TypeClass -> Flow s TypeClass)
-> Maybe TypeClass
-> Flow s TypeClass
forall b a. b -> (a -> b) -> Maybe a -> b
Y.maybe
      (String -> String -> Flow s TypeClass
forall s x. String -> String -> Flow s x
unexpected String
"type class" (String -> Flow s TypeClass) -> String -> Flow s TypeClass
forall a b. (a -> b) -> a -> b
$ Term -> String
forall a. Show a => a -> String
show Term
term) TypeClass -> Flow s TypeClass
forall a. a -> Flow s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe TypeClass -> Flow s TypeClass)
-> Maybe TypeClass -> Flow s TypeClass
forall a b. (a -> b) -> a -> b
$ Name -> Map Name TypeClass -> Maybe TypeClass
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
fn Map Name TypeClass
byName
    byName :: Map Name TypeClass
byName = [(Name, TypeClass)] -> Map Name TypeClass
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Name
_TypeClass_equality, TypeClass
TypeClassEquality), (Name
_TypeClass_ordering, TypeClass
TypeClassOrdering)]

getTypeDescription :: Type -> Flow Graph (Y.Maybe String)
getTypeDescription :: Type -> Flow Graph (Maybe String)
getTypeDescription = Map Name Term -> Flow Graph (Maybe String)
getDescription (Map Name Term -> Flow Graph (Maybe String))
-> (Type -> Map Name Term) -> Type -> Flow Graph (Maybe String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Map Name Term
typeAnnotationInternal

hasFlag :: Name -> Flow s Bool
hasFlag :: forall s. Name -> Flow s Bool
hasFlag Name
flag = Name -> Term -> Flow s Term
forall s. Name -> Term -> Flow s Term
getAttrWithDefault Name
flag (Bool -> Term
Terms.boolean Bool
False) Flow s Term -> (Term -> Flow s Bool) -> Flow s Bool
forall a b. Flow s a -> (a -> Flow s b) -> Flow s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Term -> Flow s Bool
forall s. Term -> Flow s Bool
Expect.boolean

-- | Return a zero-indexed counter for the given key: 0, 1, 2, ...
nextCount :: Name -> Flow s Int
nextCount :: forall s. Name -> Flow s Int
nextCount Name
attrName = do
  Int
count <- Name -> Term -> Flow s Term
forall s. Name -> Term -> Flow s Term
getAttrWithDefault Name
attrName (Int -> Term
Terms.int32 Int
0) Flow s Term -> (Term -> Flow s Int) -> Flow s Int
forall a b. Flow s a -> (a -> Flow s b) -> Flow s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Term -> Flow s Int
forall s. Term -> Flow s Int
Expect.int32
  Name -> Term -> Flow s ()
forall s. Name -> Term -> Flow s ()
putAttr Name
attrName (Int -> Term
Terms.int32 (Int -> Term) -> Int -> Term
forall a b. (a -> b) -> a -> b
$ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> Flow s Int
forall a. a -> Flow s a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
count

resetCount :: Name -> Flow s ()
resetCount :: forall s. Name -> Flow s ()
resetCount Name
attrName = do
  Name -> Term -> Flow s ()
forall s. Name -> Term -> Flow s ()
putAttr Name
attrName (Int -> Term
Terms.int32 Int
0)
  () -> Flow s ()
forall a. a -> Flow s a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

normalizeTermAnnotations :: Term -> Term
normalizeTermAnnotations :: Term -> Term
normalizeTermAnnotations Term
term = if Map Name Term -> Bool
forall k a. Map k a -> Bool
M.null Map Name Term
anns
    then Term
stripped
    else AnnotatedTerm -> Term
TermAnnotated (AnnotatedTerm -> Term) -> AnnotatedTerm -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Map Name Term -> AnnotatedTerm
AnnotatedTerm Term
stripped Map Name Term
anns
  where
    anns :: Map Name Term
anns = Term -> Map Name Term
termAnnotationInternal Term
term
    stripped :: Term
stripped = Term -> Term
stripTerm Term
term

normalizeTypeAnnotations :: Type -> Type
normalizeTypeAnnotations :: Type -> Type
normalizeTypeAnnotations Type
typ = if Map Name Term -> Bool
forall k a. Map k a -> Bool
M.null Map Name Term
anns
    then Type
stripped
    else AnnotatedType -> Type
TypeAnnotated (AnnotatedType -> Type) -> AnnotatedType -> Type
forall a b. (a -> b) -> a -> b
$ Type -> Map Name Term -> AnnotatedType
AnnotatedType Type
stripped Map Name Term
anns
  where
    anns :: Map Name Term
anns = Type -> Map Name Term
typeAnnotationInternal Type
typ
    stripped :: Type
stripped = Type -> Type
stripType Type
typ

putAttr :: Name -> Term -> Flow s ()
putAttr :: forall s. Name -> Term -> Flow s ()
putAttr Name
key Term
val = (s -> Trace -> FlowState s ()) -> Flow s ()
forall s x. (s -> Trace -> FlowState s x) -> Flow s x
Flow s -> Trace -> FlowState s ()
forall {s}. s -> Trace -> FlowState s ()
q
  where
    q :: s -> Trace -> FlowState s ()
q s
s0 Trace
t0 = Maybe () -> s -> Trace -> FlowState s ()
forall s x. Maybe x -> s -> Trace -> FlowState s x
FlowState (() -> Maybe ()
forall a. a -> Maybe a
Just ()) s
s0 (Trace
t0 {traceOther = M.insert key val $ traceOther t0})

setAnnotation :: Name -> Y.Maybe Term -> M.Map Name Term -> M.Map Name Term
setAnnotation :: Name -> Maybe Term -> Map Name Term -> Map Name Term
setAnnotation Name
key Maybe Term
val Map Name Term
m = (Maybe Term -> Maybe Term)
-> Name -> Map Name Term -> Map Name Term
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (Maybe Term -> Maybe Term -> Maybe Term
forall a b. a -> b -> a
const Maybe Term
val) Name
key Map Name Term
m

setDescription :: Y.Maybe String -> M.Map Name Term -> M.Map Name Term
setDescription :: Maybe String -> Map Name Term -> Map Name Term
setDescription Maybe String
d = Name -> Maybe Term -> Map Name Term -> Map Name Term
setAnnotation Name
key_description (String -> Term
Terms.string (String -> Term) -> Maybe String -> Maybe Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe String
d)

setTermAnnotation :: Name -> Y.Maybe Term -> Term -> Term
setTermAnnotation :: Name -> Maybe Term -> Term -> Term
setTermAnnotation Name
key Maybe Term
val Term
term = if Map Name Term
anns Map Name Term -> Map Name Term -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name Term
forall k a. Map k a
M.empty
    then Term
term'
    else AnnotatedTerm -> Term
TermAnnotated (AnnotatedTerm -> Term) -> AnnotatedTerm -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Map Name Term -> AnnotatedTerm
AnnotatedTerm Term
term' Map Name Term
anns
  where
    term' :: Term
term' = Term -> Term
stripTerm Term
term
    anns :: Map Name Term
anns = Name -> Maybe Term -> Map Name Term -> Map Name Term
setAnnotation Name
key Maybe Term
val (Map Name Term -> Map Name Term) -> Map Name Term -> Map Name Term
forall a b. (a -> b) -> a -> b
$ Term -> Map Name Term
termAnnotationInternal Term
term

setTermDescription :: Y.Maybe String -> Term -> Term
setTermDescription :: Maybe String -> Term -> Term
setTermDescription Maybe String
d = Name -> Maybe Term -> Term -> Term
setTermAnnotation Name
key_description (String -> Term
Terms.string (String -> Term) -> Maybe String -> Maybe Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe String
d)

-- TODO: temporary. Move this function out of Annotations
setTermType :: Y.Maybe Type -> Term -> Term
setTermType :: Maybe Type -> Term -> Term
setTermType Maybe Type
mtyp Term
term = case Maybe Type
mtyp of
    Maybe Type
Nothing -> Term -> Term
withoutType Term
term
    Just Type
typ -> TypedTerm -> Term
TermTyped (TypedTerm -> Term) -> TypedTerm -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Type -> TypedTerm
TypedTerm (Term -> Term
withoutType Term
term) Type
typ
  where
    withoutType :: Term -> Term
withoutType Term
term = case Term
term of
      TermAnnotated (AnnotatedTerm Term
term1 Map Name Term
ann) -> AnnotatedTerm -> Term
TermAnnotated (AnnotatedTerm -> Term) -> AnnotatedTerm -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Map Name Term -> AnnotatedTerm
AnnotatedTerm (Term -> Term
withoutType Term
term1) Map Name Term
ann
      TermTyped (TypedTerm Term
term1 Type
_) -> Term
term1
      Term
_ -> Term
term

setType :: Maybe Type -> Map Name Term -> Map Name Term
setType Maybe Type
mt = Name -> Maybe Term -> Map Name Term -> Map Name Term
setAnnotation Name
key_type (Type -> Term
coreEncodeType (Type -> Term) -> Maybe Type -> Maybe Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Type
mt)

setTypeAnnotation :: Name -> Y.Maybe Term -> Type -> Type
setTypeAnnotation :: Name -> Maybe Term -> Type -> Type
setTypeAnnotation Name
key Maybe Term
val Type
typ = if Map Name Term
anns Map Name Term -> Map Name Term -> Bool
forall a. Eq a => a -> a -> Bool
== Map Name Term
forall k a. Map k a
M.empty
    then Type
typ'
    else AnnotatedType -> Type
TypeAnnotated (AnnotatedType -> Type) -> AnnotatedType -> Type
forall a b. (a -> b) -> a -> b
$ Type -> Map Name Term -> AnnotatedType
AnnotatedType Type
typ' Map Name Term
anns
  where
    typ' :: Type
typ' = Type -> Type
stripType Type
typ
    anns :: Map Name Term
anns = Name -> Maybe Term -> Map Name Term -> Map Name Term
setAnnotation Name
key Maybe Term
val (Map Name Term -> Map Name Term) -> Map Name Term -> Map Name Term
forall a b. (a -> b) -> a -> b
$ Type -> Map Name Term
typeAnnotationInternal Type
typ

setTypeClasses :: M.Map Name (S.Set TypeClass) -> Type -> Type
setTypeClasses :: Map Name (Set TypeClass) -> Type -> Type
setTypeClasses Map Name (Set TypeClass)
m = Name -> Maybe Term -> Type -> Type
setTypeAnnotation Name
key_classes Maybe Term
encoded
  where
    encoded :: Maybe Term
encoded = if Map Name (Set TypeClass) -> Bool
forall k a. Map k a -> Bool
M.null Map Name (Set TypeClass)
m
      then Maybe Term
forall a. Maybe a
Nothing
      else Term -> Maybe Term
forall a. a -> Maybe a
Just (Term -> Maybe Term) -> Term -> Maybe Term
forall a b. (a -> b) -> a -> b
$ Map Term Term -> Term
Terms.map (Map Term Term -> Term) -> Map Term Term -> Term
forall a b. (a -> b) -> a -> b
$ [(Term, Term)] -> Map Term Term
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ((Name, Set TypeClass) -> (Term, Term)
encodePair ((Name, Set TypeClass) -> (Term, Term))
-> [(Name, Set TypeClass)] -> [(Term, Term)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Name (Set TypeClass) -> [(Name, Set TypeClass)]
forall k a. Map k a -> [(k, a)]
M.toList Map Name (Set TypeClass)
m)
    encodePair :: (Name, Set TypeClass) -> (Term, Term)
encodePair (Name
name, Set TypeClass
classes) = (Name -> Term
coreEncodeName Name
name, Set Term -> Term
Terms.set (Set Term -> Term) -> Set Term -> Term
forall a b. (a -> b) -> a -> b
$ [Term] -> Set Term
forall a. Ord a => [a] -> Set a
S.fromList (TypeClass -> Term
encodeClass (TypeClass -> Term) -> [TypeClass] -> [Term]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Set TypeClass -> [TypeClass]
forall a. Set a -> [a]
S.toList Set TypeClass
classes)))
    encodeClass :: TypeClass -> Term
encodeClass TypeClass
tc = Name -> Name -> Term
Terms.unitVariant Name
_TypeClass (Name -> Term) -> Name -> Term
forall a b. (a -> b) -> a -> b
$ case TypeClass
tc of
      TypeClass
TypeClassEquality -> Name
_TypeClass_equality
      TypeClass
TypeClassOrdering -> Name
_TypeClass_ordering

setTypeDescription :: Y.Maybe String -> Type -> Type
setTypeDescription :: Maybe String -> Type -> Type
setTypeDescription Maybe String
d = Name -> Maybe Term -> Type -> Type
setTypeAnnotation Name
key_description (String -> Term
Terms.string (String -> Term) -> Maybe String -> Maybe Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe String
d)

termAnnotationInternal :: Term -> M.Map Name Term
termAnnotationInternal :: Term -> Map Name Term
termAnnotationInternal = (Term -> Maybe AnnotatedTerm)
-> (AnnotatedTerm -> Term)
-> (AnnotatedTerm -> Map Name Term)
-> Term
-> Map Name Term
forall x y.
(x -> Maybe y)
-> (y -> x) -> (y -> Map Name Term) -> x -> Map Name Term
aggregateAnnotations Term -> Maybe AnnotatedTerm
getAnn AnnotatedTerm -> Term
annotatedTermSubject AnnotatedTerm -> Map Name Term
annotatedTermAnnotation
  where
    getAnn :: Term -> Maybe AnnotatedTerm
getAnn Term
t = case Term
t of
      TermAnnotated AnnotatedTerm
a -> AnnotatedTerm -> Maybe AnnotatedTerm
forall a. a -> Maybe a
Just AnnotatedTerm
a
      TermTyped (TypedTerm Term
t1 Type
_) -> Term -> Maybe AnnotatedTerm
getAnn Term
t1
      Term
_ -> Maybe AnnotatedTerm
forall a. Maybe a
Nothing

typeAnnotationInternal :: Type -> M.Map Name Term
typeAnnotationInternal :: Type -> Map Name Term
typeAnnotationInternal = (Type -> Maybe AnnotatedType)
-> (AnnotatedType -> Type)
-> (AnnotatedType -> Map Name Term)
-> Type
-> Map Name Term
forall x y.
(x -> Maybe y)
-> (y -> x) -> (y -> Map Name Term) -> x -> Map Name Term
aggregateAnnotations Type -> Maybe AnnotatedType
getAnn AnnotatedType -> Type
annotatedTypeSubject AnnotatedType -> Map Name Term
annotatedTypeAnnotation
  where
    getAnn :: Type -> Maybe AnnotatedType
getAnn Type
t = case Type
t of
      TypeAnnotated AnnotatedType
a -> AnnotatedType -> Maybe AnnotatedType
forall a. a -> Maybe a
Just AnnotatedType
a
      Type
_ -> Maybe AnnotatedType
forall a. Maybe a
Nothing

whenFlag :: Name -> Flow s a -> Flow s a -> Flow s a
whenFlag :: forall s a. Name -> Flow s a -> Flow s a -> Flow s a
whenFlag Name
flag Flow s a
fthen Flow s a
felse = do
  Bool
b <- Name -> Flow s Bool
forall s. Name -> Flow s Bool
hasFlag Name
flag
  if Bool
b
    then Flow s a
fthen
    else Flow s a
felse

-- TODO: move out of Annotations and into Rewriting
unshadowVariables :: Term -> Term
unshadowVariables :: Term -> Term
unshadowVariables Term
term = Maybe Term -> Term
forall a. HasCallStack => Maybe a -> a
Y.fromJust (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ FlowState (Set Name, Map Name Name) Term -> Maybe Term
forall s x. FlowState s x -> Maybe x
flowStateValue (FlowState (Set Name, Map Name Name) Term -> Maybe Term)
-> FlowState (Set Name, Map Name Name) Term -> Maybe Term
forall a b. (a -> b) -> a -> b
$ Flow (Set Name, Map Name Name) Term
-> (Set Name, Map Name Name)
-> Trace
-> FlowState (Set Name, Map Name Name) Term
forall s x. Flow s x -> s -> Trace -> FlowState s x
unFlow (((Term -> Flow (Set Name, Map Name Name) Term)
 -> Term -> Flow (Set Name, Map Name Name) Term)
-> Term -> Flow (Set Name, Map Name Name) Term
forall s.
((Term -> Flow s Term) -> Term -> Flow s Term)
-> Term -> Flow s Term
rewriteTermM (Term -> Flow (Set Name, Map Name Name) Term)
-> Term -> Flow (Set Name, Map Name Name) Term
rewrite Term
term) (Set Name
forall a. Set a
S.empty, Map Name Name
forall k a. Map k a
M.empty) Trace
emptyTrace
  where
    rewrite :: (Term -> Flow (Set Name, Map Name Name) Term)
-> Term -> Flow (Set Name, Map Name Name) Term
rewrite Term -> Flow (Set Name, Map Name Name) Term
recurse Term
term = do
      (Set Name
reserved, Map Name Name
subst) <- Flow (Set Name, Map Name Name) (Set Name, Map Name Name)
forall s. Flow s s
getState
      case Term
term of
        TermVariable Name
v -> Term -> Flow (Set Name, Map Name Name) Term
forall a. a -> Flow (Set Name, Map Name Name) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Term -> Flow (Set Name, Map Name Name) Term)
-> Term -> Flow (Set Name, Map Name Name) Term
forall a b. (a -> b) -> a -> b
$ Name -> Term
TermVariable (Name -> Term) -> Name -> Term
forall a b. (a -> b) -> a -> b
$ Name -> Maybe Name -> Name
forall a. a -> Maybe a -> a
Y.fromMaybe Name
v (Maybe Name -> Name) -> Maybe Name -> Name
forall a b. (a -> b) -> a -> b
$ Name -> Map Name Name -> Maybe Name
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Name
v Map Name Name
subst
        TermFunction (FunctionLambda (Lambda Name
v Maybe Type
d Term
body)) -> if Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
S.member Name
v Set Name
reserved
          then do
            Name
v' <- Flow (Set Name, Map Name Name) Name
forall {s}. Flow s Name
freshName
            (Set Name, Map Name Name) -> Flow (Set Name, Map Name Name) ()
forall s. s -> Flow s ()
putState (Name -> Set Name -> Set Name
forall a. Ord a => a -> Set a -> Set a
S.insert Name
v' Set Name
reserved, Name -> Name -> Map Name Name -> Map Name Name
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Name
v Name
v' Map Name Name
subst)
            Term
body' <- Term -> Flow (Set Name, Map Name Name) Term
recurse Term
body
            (Set Name, Map Name Name) -> Flow (Set Name, Map Name Name) ()
forall s. s -> Flow s ()
putState (Set Name
reserved, Map Name Name
subst)
            Term -> Flow (Set Name, Map Name Name) Term
forall a. a -> Flow (Set Name, Map Name Name) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Term -> Flow (Set Name, Map Name Name) Term)
-> Term -> Flow (Set Name, Map Name Name) Term
forall a b. (a -> b) -> a -> b
$ Function -> Term
TermFunction (Function -> Term) -> Function -> Term
forall a b. (a -> b) -> a -> b
$ Lambda -> Function
FunctionLambda (Lambda -> Function) -> Lambda -> Function
forall a b. (a -> b) -> a -> b
$ Name -> Maybe Type -> Term -> Lambda
Lambda Name
v' Maybe Type
d Term
body'
          else do
            (Set Name, Map Name Name) -> Flow (Set Name, Map Name Name) ()
forall s. s -> Flow s ()
putState (Name -> Set Name -> Set Name
forall a. Ord a => a -> Set a -> Set a
S.insert Name
v Set Name
reserved, Map Name Name
subst)
            Term
body' <- Term -> Flow (Set Name, Map Name Name) Term
recurse Term
body
            Term -> Flow (Set Name, Map Name Name) Term
forall a. a -> Flow (Set Name, Map Name Name) a
forall (m :: * -> *) a. Monad m => a -> m a
return (Term -> Flow (Set Name, Map Name Name) Term)
-> Term -> Flow (Set Name, Map Name Name) Term
forall a b. (a -> b) -> a -> b
$ Function -> Term
TermFunction (Function -> Term) -> Function -> Term
forall a b. (a -> b) -> a -> b
$ Lambda -> Function
FunctionLambda (Lambda -> Function) -> Lambda -> Function
forall a b. (a -> b) -> a -> b
$ Name -> Maybe Type -> Term -> Lambda
Lambda Name
v Maybe Type
d Term
body'
        Term
_ -> Term -> Flow (Set Name, Map Name Name) Term
recurse Term
term
    freshName :: Flow s Name
freshName = (\Int
n -> String -> Name
Name (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"s" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n) (Int -> Name) -> Flow s Int -> Flow s Name
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Flow s Int
forall s. Name -> Flow s Int
nextCount (String -> Name
Name String
"unshadow")