1 files changed, 13 insertions, 0 deletions
diff --git a/src/CHAD/Analysis/Identity.hs b/src/CHAD/Analysis/Identity.hs
index 212cc7d..284ab49 100644
--- a/src/CHAD/Analysis/Identity.hs
+++ b/src/CHAD/Analysis/Identity.hs
@@ -34,6 +34,7 @@ data ValId t where
   VIArr :: Int -> Vec n Int -> ValId (TArr n t)
   VIScal :: Int -> ValId (TScal t)
   VIAccum :: Int -> ValId (TAccum t)
+  VIUser :: ValId (UserRep t) -> ValId (TUser t)
 deriving instance Show (ValId t)
 
 instance PrettyX ValId where
@@ -56,6 +57,7 @@ instance PrettyX ValId where
     VIArr i is -> 'A' : show i ++ "[" ++ intercalate "," (map show (toList is)) ++ "]"
     VIScal i -> show i
     VIAccum i -> 'C' : show i
+    VIUser a -> 'U' : show a
 
 validSplitEither :: ValId (TEither a b) -> (Maybe (ValId a), Maybe (ValId b))
 validSplitEither (VIEither (Left v)) = (Just v, Nothing)
@@ -386,6 +388,15 @@ idana env expr = case expr of
     res <- genIds t
     pure (res, EError res t s)
 
+  EUser _ t e -> do
+    (v, e') <- idana env e
+    pure (VIUser v, EUser (VIUser v) t e')
+
+  EUnUser _ e -> do
+    (v, e') <- idana env e
+    let VIUser v' = v
+    pure (v', EUnUser v' e')
+
 -- | This value might be either of the two arguments; we don't know which.
 unify :: ValId t -> ValId t -> IdGen (ValId t)
 unify VINil VINil = pure VINil
@@ -412,6 +423,7 @@ unify (VILEither a) (VILEither b) = VILEither <$> unify a b
 unify (VIArr i is) (VIArr j js) = VIArr <$> unifyID i j <*> vecZipWithA unifyID is js
 unify (VIScal i) (VIScal j) = VIScal <$> unifyID i j
 unify (VIAccum i) (VIAccum j) = VIAccum <$> unifyID i j
+unify (VIUser i) (VIUser j) = VIUser <$> unify i j
 
 unifyID :: Int -> Int -> IdGen Int
 unifyID i j | i == j = pure i
@@ -426,6 +438,7 @@ genIds (STMaybe t) = VIMaybe' <$> genIds t
 genIds (STArr n _) = VIArr <$> genId <*> vecReplicateA n genId
 genIds STScal{} = VIScal <$> genId
 genIds STAccum{} = VIAccum <$> genId
+genIds (STUser t) = VIUser <$> genIds (userRepTy t)
 
 shidsToVec :: SNat n -> ValId (Tup (Replicate n TIx)) -> IdGen (Vec n Int)
 shidsToVec SZ _ = pure VNil