diff options
Diffstat (limited to 'src/AST/Count.hs')
| -rw-r--r-- | src/AST/Count.hs | 21 |
1 files changed, 11 insertions, 10 deletions
diff --git a/src/AST/Count.hs b/src/AST/Count.hs index 7e70a7d..289c1fb 100644 --- a/src/AST/Count.hs +++ b/src/AST/Count.hs @@ -76,17 +76,17 @@ scaleManyOccEnv :: OccEnv env -> OccEnv env scaleManyOccEnv OccEnd = OccEnd scaleManyOccEnv (OccPush e o) = OccPush (scaleManyOccEnv e) (scaleMany o) +occEnvPop :: OccEnv (t : env) -> OccEnv env +occEnvPop (OccPush o _) = o +occEnvPop OccEnd = OccEnd + occCountAll :: Expr x env t -> OccEnv env -occCountAll = occCountGeneral onehotOccEnv unpush unpushN (<||>!) scaleManyOccEnv +occCountAll = occCountGeneral onehotOccEnv occEnvPop occEnvPopN (<||>!) scaleManyOccEnv where - unpush :: OccEnv (t : env) -> OccEnv env - unpush (OccPush o _) = o - unpush OccEnd = OccEnd - - unpushN :: SNat n -> OccEnv (ConsN n TIx env) -> OccEnv env - unpushN _ OccEnd = OccEnd - unpushN SZ e = e - unpushN (SS n) (OccPush e _) = unpushN n e + occEnvPopN :: SNat n -> OccEnv (ConsN n TIx env) -> OccEnv env + occEnvPopN _ OccEnd = OccEnd + occEnvPopN SZ e = e + occEnvPopN (SS n) (OccPush e _) = occEnvPopN n e occCountGeneral :: forall r env t x. (forall env'. Monoid (r env')) @@ -112,11 +112,12 @@ occCountGeneral onehot unpush unpushN alter many = go EBuild1 _ a b -> go a <> many (unpush (go b)) EBuild _ es e -> foldMap go es <> many (unpushN (vecLength es) (go e)) EFold1 _ a b -> many (unpush (unpush (go a))) <> go b + EUnit _ e -> go e EConst{} -> mempty EIdx0 _ e -> go e EIdx1 _ a b -> go a <> go b EIdx _ e es -> go e <> foldMap go es EOp _ _ e -> go e EWith a b -> go a <> unpush (go b) - EAccum a b e -> go a <> go b <> go e + EAccum1 a b e -> go a <> go b <> go e EError{} -> mempty |