-- | 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 :: String
key_classes = String
"classes" :: String
key_description :: String
key_description = String
"description"
key_type :: String
key_type = String
"type"


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

failOnFlag :: String -> String -> Flow s ()
failOnFlag :: forall s. String -> String -> Flow s ()
failOnFlag String
flag String
msg = do
  Bool
val <- String -> Flow s Bool
forall s. String -> Flow s Bool
hasFlag String
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 :: String -> Flow s (Maybe Term)
getAttr :: forall s. String -> Flow s (Maybe Term)
getAttr String
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
$ String -> Map String Term -> Maybe Term
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup String
key (Map String Term -> Maybe Term) -> Map String Term -> Maybe Term
forall a b. (a -> b) -> a -> b
$ Trace -> Map String Term
traceOther Trace
t0) s
s0 Trace
t0

getAttrWithDefault :: String -> Term -> Flow s Term
getAttrWithDefault :: forall s. String -> Term -> Flow s Term
getAttrWithDefault String
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
<$> String -> Flow s (Maybe Term)
forall s. String -> Flow s (Maybe Term)
getAttr String
key

getDescription :: M.Map String Term -> Flow Graph (Y.Maybe String)
getDescription :: Map String Term -> Flow Graph (Maybe String)
getDescription Map String Term
anns = case String -> Map String Term -> Maybe Term
getAnnotation String
key_description Map String 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]
++ String -> String
forall a. Show a => a -> String
show String
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 :: String -> Term -> Y.Maybe Term
getTermAnnotation :: String -> Term -> Maybe Term
getTermAnnotation String
key = String -> Map String Term -> Maybe Term
getAnnotation String
key (Map String Term -> Maybe Term)
-> (Term -> Map String Term) -> Term -> Maybe Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term -> Map String Term
termAnnotationInternal

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

getType :: M.Map String Term -> Flow Graph (Y.Maybe Type)
getType :: Map String Term -> Flow Graph (Maybe Type)
getType Map String Term
anns = case String -> Map String Term -> Maybe Term
getAnnotation String
key_type Map String 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 :: String -> Type -> Y.Maybe Term
getTypeAnnotation :: String -> Type -> Maybe Term
getTypeAnnotation String
key = String -> Map String Term -> Maybe Term
getAnnotation String
key (Map String Term -> Maybe Term)
-> (Type -> Map String Term) -> Type -> Maybe Term
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Map String Term
typeAnnotationInternal

getTypeClasses :: Type -> Flow Graph (M.Map Name (S.Set TypeClass))
getTypeClasses :: Type -> Flow Graph (Map Name (Set TypeClass))
getTypeClasses Type
typ = case String -> Type -> Maybe Term
getTypeAnnotation String
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 String Term -> Flow Graph (Maybe String)
getDescription (Map String Term -> Flow Graph (Maybe String))
-> (Type -> Map String Term) -> Type -> Flow Graph (Maybe String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> Map String Term
typeAnnotationInternal

hasFlag :: String -> Flow s Bool
hasFlag :: forall s. String -> Flow s Bool
hasFlag String
flag = String -> Term -> Flow s Term
forall s. String -> Term -> Flow s Term
getAttrWithDefault String
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 :: String -> Flow s Int
nextCount :: forall s. String -> Flow s Int
nextCount String
attrName = do
  Int
count <- String -> Term -> Flow s Term
forall s. String -> Term -> Flow s Term
getAttrWithDefault String
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
  String -> Term -> Flow s ()
forall s. String -> Term -> Flow s ()
putAttr String
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 :: String -> Flow s ()
resetCount :: forall s. String -> Flow s ()
resetCount String
attrName = do
  String -> Term -> Flow s ()
forall s. String -> Term -> Flow s ()
putAttr String
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 String Term -> Bool
forall k a. Map k a -> Bool
M.null Map String Term
anns
    then Term
stripped
    else AnnotatedTerm -> Term
TermAnnotated (AnnotatedTerm -> Term) -> AnnotatedTerm -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Map String Term -> AnnotatedTerm
AnnotatedTerm Term
stripped Map String Term
anns
  where
    anns :: Map String Term
anns = Term -> Map String Term
termAnnotationInternal Term
term
    stripped :: Term
stripped = Term -> Term
stripTerm Term
term

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

putAttr :: String -> Term -> Flow s ()
putAttr :: forall s. String -> Term -> Flow s ()
putAttr String
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 :: String -> Y.Maybe Term -> M.Map String Term -> M.Map String Term
setAnnotation :: String -> Maybe Term -> Map String Term -> Map String Term
setAnnotation String
key Maybe Term
val Map String Term
m = (Maybe Term -> Maybe Term)
-> String -> Map String Term -> Map String 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) String
key Map String Term
m

setDescription :: Y.Maybe String -> M.Map String Term -> M.Map String Term
setDescription :: Maybe String -> Map String Term -> Map String Term
setDescription Maybe String
d = String -> Maybe Term -> Map String Term -> Map String Term
setAnnotation String
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 :: String -> Y.Maybe Term -> Term -> Term
setTermAnnotation :: String -> Maybe Term -> Term -> Term
setTermAnnotation String
key Maybe Term
val Term
term = if Map String Term
anns Map String Term -> Map String Term -> Bool
forall a. Eq a => a -> a -> Bool
== Map String 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 String Term -> AnnotatedTerm
AnnotatedTerm Term
term' Map String Term
anns
  where
    term' :: Term
term' = Term -> Term
stripTerm Term
term
    anns :: Map String Term
anns = String -> Maybe Term -> Map String Term -> Map String Term
setAnnotation String
key Maybe Term
val (Map String Term -> Map String Term)
-> Map String Term -> Map String Term
forall a b. (a -> b) -> a -> b
$ Term -> Map String Term
termAnnotationInternal Term
term

setTermDescription :: Y.Maybe String -> Term -> Term
setTermDescription :: Maybe String -> Term -> Term
setTermDescription Maybe String
d = String -> Maybe Term -> Term -> Term
setTermAnnotation String
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 String Term
ann) -> AnnotatedTerm -> Term
TermAnnotated (AnnotatedTerm -> Term) -> AnnotatedTerm -> Term
forall a b. (a -> b) -> a -> b
$ Term -> Map String Term -> AnnotatedTerm
AnnotatedTerm (Term -> Term
withoutType Term
term1) Map String Term
ann
      TermTyped (TypedTerm Term
term1 Type
_) -> Term
term1
      Term
_ -> Term
term

setType :: Maybe Type -> Map String Term -> Map String Term
setType Maybe Type
mt = String -> Maybe Term -> Map String Term -> Map String Term
setAnnotation String
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 :: String -> Y.Maybe Term -> Type -> Type
setTypeAnnotation :: String -> Maybe Term -> Type -> Type
setTypeAnnotation String
key Maybe Term
val Type
typ = if Map String Term
anns Map String Term -> Map String Term -> Bool
forall a. Eq a => a -> a -> Bool
== Map String 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 String Term -> AnnotatedType
AnnotatedType Type
typ' Map String Term
anns
  where
    typ' :: Type
typ' = Type -> Type
stripType Type
typ
    anns :: Map String Term
anns = String -> Maybe Term -> Map String Term -> Map String Term
setAnnotation String
key Maybe Term
val (Map String Term -> Map String Term)
-> Map String Term -> Map String Term
forall a b. (a -> b) -> a -> b
$ Type -> Map String 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 = String -> Maybe Term -> Type -> Type
setTypeAnnotation String
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 = String -> Maybe Term -> Type -> Type
setTypeAnnotation String
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 String Term
termAnnotationInternal :: Term -> Map String Term
termAnnotationInternal = (Term -> Maybe AnnotatedTerm)
-> (AnnotatedTerm -> Term)
-> (AnnotatedTerm -> Map String Term)
-> Term
-> Map String Term
forall x y.
(x -> Maybe y)
-> (y -> x) -> (y -> Map String Term) -> x -> Map String Term
aggregateAnnotations Term -> Maybe AnnotatedTerm
getAnn AnnotatedTerm -> Term
annotatedTermSubject AnnotatedTerm -> Map String 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 String Term
typeAnnotationInternal :: Type -> Map String Term
typeAnnotationInternal = (Type -> Maybe AnnotatedType)
-> (AnnotatedType -> Type)
-> (AnnotatedType -> Map String Term)
-> Type
-> Map String Term
forall x y.
(x -> Maybe y)
-> (y -> x) -> (y -> Map String Term) -> x -> Map String Term
aggregateAnnotations Type -> Maybe AnnotatedType
getAnn AnnotatedType -> Type
annotatedTypeSubject AnnotatedType -> Map String 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 :: String -> Flow s a -> Flow s a -> Flow s a
whenFlag :: forall s a. String -> Flow s a -> Flow s a -> Flow s a
whenFlag String
flag Flow s a
fthen Flow s a
felse = do
  Bool
b <- String -> Flow s Bool
forall s. String -> Flow s Bool
hasFlag String
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)
-> (Map String Term
    -> Flow (Set Name, Map Name Name) (Map String Term))
-> Term
-> Flow (Set Name, Map Name Name) Term
forall s.
((Term -> Flow s Term) -> Term -> Flow s Term)
-> (Map String Term -> Flow s (Map String Term))
-> Term
-> Flow s Term
rewriteTermM (Term -> Flow (Set Name, Map Name Name) Term)
-> Term -> Flow (Set Name, Map Name Name) Term
rewrite (Map String Term -> Flow (Set Name, Map Name Name) (Map String Term)
forall a. a -> Flow (Set Name, Map Name Name) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map String Term
 -> Flow (Set Name, Map Name Name) (Map String Term))
-> (Map String Term -> Map String Term)
-> Map String Term
-> Flow (Set Name, Map Name Name) (Map String Term)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map String Term -> Map String Term
forall a. a -> a
id) 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 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 -> Term -> Lambda
Lambda Name
v' 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 -> Term -> Lambda
Lambda Name
v 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
<$> String -> Flow s Int
forall s. String -> Flow s Int
nextCount String
"unshadow"