WIP identity analysis

3 files changed, 280 insertions, 0 deletions

diff --git a/chad-fast.cabal b/chad-fast.cabal
index 6635f6c..63479bb 100644
--- a/chad-fast.cabal
+++ b/chad-fast.cabal
@@ -10,6 +10,7 @@ build-type:          Simple
 
 library
   exposed-modules:
+    Analysis.Identity
     Array
     AST
     AST.Bindings
@@ -43,6 +44,7 @@ library
     Lemmas
     -- PreludeCu
     Simplify
+    Util.IdGen
   other-modules:
   build-depends:
     base >= 4.19 && < 4.21,
diff --git a/src/Analysis/Identity.hs b/src/Analysis/Identity.hs
new file mode 100644
index 0000000..5c398d2
--- /dev/null
+++ b/src/Analysis/Identity.hs
@@ -0,0 +1,260 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE GADTs #-}
+module Analysis.Identity (
+  identityAnalysis,
+) where
+
+import Data.Functor.Const
+
+import AST
+import Data
+import Util.IdGen
+
+
+-- | Every array, scalar and accumulator has an ID. Trivial values such as
+-- Nothing only have the knowledge that they are indeed Nothing. Compound
+-- values know which values they consist of.
+data ValId t where
+  VINil :: ValId TNil
+  VIPair :: ValId a -> ValId b -> ValId (TPair a b)
+  VIEither :: Either (ValId a) (ValId b) -> ValId (TEither a b)  -- ^ known alternative
+  VIEither' :: ValId a -> ValId b -> ValId (TEither a b)  -- ^ unknown alternative, but known values in each case
+  VIMaybe :: Maybe (ValId a) -> ValId (TMaybe a)
+  VIArr :: Int -> ValId (TArr n t)
+  VIScal :: Int -> ValId (TScal t)
+  VIAccum :: Int -> ValId (TAccum t)
+
+  -- | We don't know what this consists of, but it's a value, and let's just
+  -- give it an ID nevertheless.
+  VIThing :: Int -> ValId t
+
+instance Eq (ValId t) where
+  VINil         == VINil           = True
+  VINil         == _               = False
+  VIPair a b    == VIPair a' b'    = a == a' && b == b'
+  VIPair{}      == _               = False
+  VIEither a    == VIEither a'     = a == a'
+  VIEither{}    == _               = False
+  VIEither' a b == VIEither' a' b' = a == a' && b == b'
+  VIEither'{}   == _               = False
+  VIMaybe a     == VIMaybe a'      = a == a'
+  VIMaybe{}     == _               = False
+  VIArr i       == VIArr i'        = i == i'
+  VIArr{}       == _               = False
+  VIScal i      == VIScal i'       = i == i'
+  VIScal{}      == _               = False
+  VIAccum i     == VIAccum i'      = i == i'
+  VIAccum{}     == _               = False
+  VIThing i     == VIThing i'      = i == i'
+  VIThing{}     == _               = False
+
+-- | Symbolic partial evaluation.
+identityAnalysis :: SList STy env -> Expr x env t -> Expr ValId env t
+identityAnalysis env term = runIdGen 0 $ do
+  env' <- slistMapA numberConstant env
+  snd <$> idana env' term
+  where
+    numberConstant :: STy t -> IdGen (ValId t)
+    numberConstant = \case
+      STNil -> pure VINil
+      STPair a b -> VIPair <$> numberConstant a <*> numberConstant b
+      STEither a b -> VIEither' <$> numberConstant a <*> numberConstant b
+      STMaybe{} -> VIThing <$> genId
+      STArr{} -> VIArr <$> genId
+      STScal{} -> VIScal <$> genId
+      STAccum{} -> VIAccum <$> genId
+
+idana :: SList ValId env -> Expr x env t -> IdGen (ValId t, Expr ValId env t)
+idana env = \case
+  EVar _ t i -> do
+    let v = slistIdx env i
+    pure (v, EVar v t i)
+
+  ELet _ e1 e2 -> do
+    (v1, e1') <- idana env e1
+    (v2, e2') <- idana (v1 `SCons` env) e2
+    pure (v2, ELet v2 e1' e2')
+
+  EPair _ e1 e2 -> do
+    (v1, e1') <- idana env e1
+    (v2, e2') <- idana env e2
+    pure (VIPair v1 v2, EPair (VIPair v1 v2) e1' e2')
+
+  EFst _ e -> do
+    (v, e') <- idana env e
+    v' <- case v of VIPair v1 _ -> pure v1
+                    _ -> VIThing <$> genId
+    pure (v', EFst v' e')
+
+  ESnd _ e -> do
+    (v, e') <- idana env e
+    v' <- case v of VIPair _ v2 -> pure v2
+                    _ -> VIThing <$> genId
+    pure (v', ESnd v' e')
+
+  ENil _ -> pure (VINil, ENil VINil)
+
+  EInl _ t2 e1 -> do
+    (v1, e1') <- idana env e1
+    let v = VIEither (Left v1)
+    pure (v, EInl v t2 e1')
+
+  EInr _ t1 e2 -> do
+    (v2, e2') <- idana env e2
+    let v = VIEither (Right v2)
+    pure (v, EInr v t1 e2')
+
+  ECase _ e1 e2 e3 -> do
+    (v1, e1') <- idana env e1
+    case v1 of
+      VIEither (Left v1') -> do
+        (v2, e2') <- idana (v1' `SCons` env) e2
+        scrap <- VIThing <$> genId
+        (_, e3') <- idana (scrap `SCons` env) e3
+        pure (v2, ECase v2 e1' e2' e3')
+      VIEither (Right v1') -> do
+        scrap <- VIThing <$> genId
+        (_, e2') <- idana (scrap `SCons` env) e2
+        (v3, e3') <- idana (v1' `SCons` env) e3
+        pure (v3, ECase v3 e1' e2' e3')
+      VIEither' v1'l v1'r -> do
+        (_, e2') <- idana (v1'l `SCons` env) e2
+        (_, e3') <- idana (v1'r `SCons` env) e3
+        res <- genId
+        pure (VIThing res, ECase (VIThing res) e1' e2' e3')
+      VIThing _ -> do
+        x2 <- genId
+        x3 <- genId
+        (v2, e2') <- idana (VIThing x2 `SCons` env) e2
+        (v3, e3') <- idana (VIThing x3 `SCons` env) e3
+        if v2 == v3
+          then pure (v2, ECase v2 e1' e2' e3')
+          else do
+            res <- genId
+            pure (VIThing res, ECase (VIThing res) e1' e2' e3')
+
+  ENothing _ t -> pure (VIMaybe Nothing, ENothing (VIMaybe Nothing) t)
+
+  EJust _ e1 -> do
+    (v1, e1') <- idana env e1
+    let v = VIMaybe (Just v1)
+    pure (v, EJust v e1')
+
+  EMaybe _ e1 e2 e3 -> do
+    (v3, e3') <- idana env e3
+    case v3 of
+      VIMaybe Nothing -> do
+        (v1, e1') <- idana env e1
+        scrap <- VIThing <$> genId
+        (_, e2') <- idana (scrap `SCons` env) e2
+        pure (v1, EMaybe v1 e1' e2' e3')
+      VIMaybe (Just v3j) -> do
+        (v2, e2') <- idana (v3j `SCons` env) e2
+        (_, e1') <- idana env e1
+        pure (v2, EMaybe v2 e1' e2' e3')
+      VIThing _ -> do
+        (v1, e1') <- idana env e1
+        scrap <- VIThing <$> genId
+        (v2, e2') <- idana (scrap `SCons` env) e2
+        if v1 == v2
+          then pure (v2, EMaybe v2 e1' e2' e3')
+          else do
+            res <- genId
+            pure (VIThing res, EMaybe (VIThing res) e1' e2' e3')
+
+  EConstArr _ dim t arr -> do
+    x1 <- VIArr <$> genId
+    pure (x1, EConstArr x1 dim t arr)
+
+  EBuild _ dim e1 e2 -> do
+    (_, e1') <- idana env e1
+    scrap <- VIThing <$> genId
+    (_, e2') <- idana (scrap `SCons` env) e2
+    res <- VIArr <$> genId
+    pure (res, EBuild res dim e1' e2')
+
+  EFold1Inner _ e1 e2 e3 -> do
+    scrap1 <- VIThing <$> genId
+    scrap2 <- VIThing <$> genId
+    (_, e1') <- idana (scrap1 `SCons` scrap2 `SCons` env) e1
+    (_, e2') <- idana env e2
+    (_, e3') <- idana env e3
+    res <- VIArr <$> genId
+    pure (res, EFold1Inner res e1' e2' e3')
+
+  ESum1Inner _ e1 -> do
+    (_, e1') <- idana env e1
+    res <- VIArr <$> genId
+    pure (res, ESum1Inner res e1')
+
+  EUnit _ e1 -> do
+    (_, e1') <- idana env e1
+    res <- VIArr <$> genId
+    pure (res, EUnit res e1')
+
+  EReplicate1Inner _ e1 e2 -> do
+    (_, e1') <- idana env e1
+    (_, e2') <- idana env e2
+    res <- VIArr <$> genId
+    pure (res, EReplicate1Inner res e1' e2')
+
+  EMaximum1Inner _ e1 -> do
+    (_, e1') <- idana env e1
+    res <- VIArr <$> genId
+    pure (res, EMaximum1Inner res e1')
+
+  EMinimum1Inner _ e1 -> do
+    (_, e1') <- idana env e1
+    res <- VIArr <$> genId
+    pure (res, EMinimum1Inner res e1')
+
+  EConst _ t val -> do
+    res <- VIScal <$> genId
+    pure (res, EConst res t val)
+
+  EIdx0 _ e1 -> do
+    (_, e1') <- idana env e1
+    res <- VIThing <$> genId
+    pure (res, EIdx0 res e1')
+
+  EIdx1 _ e1 e2 -> do
+    (_, e1') <- idana env e1
+    (_, e2') <- idana env e2
+    res <- VIThing <$> genId
+    pure (res, EIdx1 res e1' e2')
+
+  EIdx _ e1 e2 -> do
+    (_, e1') <- idana env e1
+    (_, e2') <- idana env e2
+    res <- VIThing <$> genId
+    pure (res, EIdx res e1' e2')
+
+  EShape _ e1 -> do
+    (_, e1') <- idana env e1
+    res <- VIThing <$> genId
+    pure (res, EShape res e1')
+
+  EOp _ op e1 -> do
+    (_, e1') <- idana env e1
+    res <- VIThing <$> genId
+    pure (res, EOp res op e1')
+
+  ECustom _ t1 t2 t3 e1 e2 e3 e4 e5 -> do
+    x1 <- VIThing <$> genId
+    x2 <- VIThing <$> genId
+    (_, e1') <- idana (x1 `SCons` x2 `SCons` SNil) e1
+    x3 <- VIThing <$> genId
+    x4 <- VIThing <$> genId
+    (_, e2') <- idana (x3 `SCons` x4 `SCons` SNil) e2
+    x5 <- VIThing <$> genId
+    x6 <- VIThing <$> genId
+    (_, e3') <- idana (x5 `SCons` x6 `SCons` SNil) e3
+    (_, e4') <- idana env e4
+    (_, e5') <- idana env e5
+    res <- VIThing <$> genId
+    pure (res, ECustom res t1 t2 t3 e1' e2' e3' e4' e5')
+
+  -- EWith e1 e2
+
+  _ -> _
diff --git a/src/Util/IdGen.hs b/src/Util/IdGen.hs
new file mode 100644
index 0000000..fcfb6e7
--- /dev/null
+++ b/src/Util/IdGen.hs
@@ -0,0 +1,18 @@
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+module Util.IdGen where
+
+import Control.Monad.Trans.State.Strict
+
+
+newtype IdGen a = IdGen (State Int a)
+  deriving newtype (Functor, Applicative, Monad)
+
+genId :: IdGen Int
+genId = IdGen (state (\i -> (i, i + 1)))
+
+runIdGen :: Int -> IdGen a -> a
+runIdGen start (IdGen m) = evalState m start
+
+runIdGen' :: Int -> IdGen a -> (a, Int)
+runIdGen' start (IdGen m) = runState m start