handle vector operations with combinators

Author: martyall

bundleUnbundle.dot

@@ -12,7 +12,7 @@ source-repository-package
 source-repository-package
   type: git
   location: https://github.com/l-adic/galois-fields.git
-  tag: fc82039e811ba68c10527cf871796b7ac8514926
+  tag: b0867ffdebda5043c80315a51b15e82ed25acba6
   --sha256: j/zGFd2aeowzJfgCCBmJYmG8mDsfF0irqj/cPOw9ulE=
 
 source-repository-package

cabal.project

@@ -47,11 +47,9 @@ instance Bifunctor AffineCircuit where
     ScalarMul s x -> ScalarMul (f s) (bimap f g x)
     ConstGate c -> ConstGate (f c)
     Var i -> Var (g i)
-
     Nil -> Nil
 
 instance (Pretty i, Pretty f) => Pretty (AffineCircuit f i) where
-  pretty :: AffineCircuit f i -> Doc
   pretty = prettyPrec 0
     where
       prettyPrec :: Int -> AffineCircuit f i -> Doc
@@ -61,8 +59,7 @@ instance (Pretty i, Pretty f) => Pretty (AffineCircuit f i) where
             text "nil"
           Var v ->
             pretty v
-          ConstGate f ->
-            pretty f
+          ConstGate f -> pretty f
           ScalarMul f e1 ->
             pretty f <+> text "*" <+> parensPrec 7 p (prettyPrec p e1)
           Add e1 e2 ->

circuit/src/Circuit/Affine.hs

@@ -16,7 +16,7 @@ module Circuit.Arithmetic
     CircuitVars (..),
     relabel,
     collectCircuitVars,
-    booleanWires
+    booleanWires,
   )
 where
 
@@ -342,4 +342,4 @@ booleanWires :: ArithCircuit f -> Set Wire
 booleanWires (ArithCircuit gates) = foldMap f gates
   where
     f (Boolean i) = Set.singleton i
-    f _ = mempty
+    f _ = mempty

circuit/src/Circuit/Arithmetic.hs

@@ -7,9 +7,9 @@ where
 
 import Circuit.Affine ()
 import Circuit.Arithmetic (ArithCircuit (..), Gate (..), Wire (..), booleanWires)
+import Data.Field.Galois (PrimeField)
 import Data.Text qualified as Text
 import Protolude
-import Data.Field.Galois (PrimeField)
 import System.FilePath (replaceExtension)
 import System.Process (readProcessWithExitCode)
 import Text.PrettyPrint.Leijen.Text (Pretty (..))
@@ -30,7 +30,7 @@ arithCircuitToDot c@(ArithCircuit gates) =
     dotArrow s t = s <> " -> " <> t
 
     dotArrowLabel :: Text -> Text -> Text -> Text -> Text
-    dotArrowLabel s t lbl _c = 
+    dotArrowLabel s t lbl _c =
       dotArrow s t <> " [label=\"" <> lbl <> "\", color=\"" <> _c <> "\"]"
 
     labelNode lblId lbl = lblId <> " [label=\"" <> lbl <> "\"]"
@@ -53,8 +53,10 @@ arithCircuitToDot c@(ArithCircuit gates) =
         gateLabel = dotWire output
         inputs circuit tgt =
           map
-            ( \a -> (\src -> dotArrowLabel src tgt (show $ pretty src) (color a))
-                . dotWire $ a
+            ( \a ->
+                (\src -> dotArrowLabel src tgt (show $ pretty src) (color a))
+                  . dotWire
+                  $ a
             )
             $ toList circuit
     graphGate (Equal i m output) =
@@ -73,8 +75,8 @@ arithCircuitToDot c@(ArithCircuit gates) =
       where
         gateLabel = Text.concat . fmap dotWire $ outputs
     graphGate (Boolean _) = []
-    
-    -- gateLabel = dotWire i
+
+-- gateLabel = dotWire i
 
 callDot :: Text -> IO Text
 callDot g = do

circuit/src/Circuit/Dot.hs

@@ -12,7 +12,7 @@ module Circuit.Language.Compile
     compileWithWire,
     compileWithWires,
     exprToArithCircuit,
-    unBundle
+    _unBundle,
   )
 where
 
@@ -21,21 +21,21 @@ import Circuit.Arithmetic
 import Circuit.Language.Expr
   ( BinOp (..),
     Expr (..),
+    Hash (Hash),
     UVar (..),
     UnOp (..),
     getAnnotation,
     hashCons,
-    Hash (Hash),
   )
 import Circuit.Language.TExpr qualified as TExpr
 import Data.Field.Galois (GaloisField)
 import Data.Map qualified as Map
+import Data.Maybe (fromJust)
 import Data.Set qualified as Set
 import Data.Vector qualified as V
+import Data.Vector.Sized qualified as SV
 import Protolude hiding (Semiring)
 import Text.PrettyPrint.Leijen.Text hiding ((<$>))
-import Data.Maybe (fromJust)
-import Data.Vector.Sized qualified as SV
 import Unsafe.Coerce (unsafeCoerce)
 
 -------------------------------------------------------------------------------
@@ -204,13 +204,13 @@ compileWithWire freshWire e = do
     <$> compileWithWires (V.singleton $ fmap TExpr.coerceGroundType freshWire) e
 
 compileWithWires ::
-  (Hashable f) => 
+  (Hashable f) =>
   (GaloisField f) =>
   (MonadState (BuilderState f) m) =>
   (MonadError (CircuitBuilderError f) m) =>
   V.Vector (m (TExpr.Var Wire f f)) ->
   TExpr.Expr Wire f ty ->
-  m (V.Vector (TExpr.Var Wire f f) )
+  m (V.Vector (TExpr.Var Wire f f))
 compileWithWires ws expr = do
   e <- hashCons <$> unType expr
   compileOut <- memoizedCompile e
@@ -331,7 +331,7 @@ _compile expr = withCompilerCache (getAnnotation expr) $ case expr of
   ESplit _ n input -> do
     -- assertSingle is justified as the input must be of type f
     i <- memoizedCompile input >>= assertSingleSource >>= addWire
-    outputs <- V.generateM n $ \_ -> do 
+    outputs <- V.generateM n $ \_ -> do
       w <- imm
       emit $ Boolean w
       pure w
@@ -351,14 +351,6 @@ _compile expr = withCompilerCache (getAnnotation expr) $ case expr of
     bs <- toList <$> memoizedCompile bits
     ws <- traverse addWire bs
     pure . V.singleton . AffineSource $ unsplit ws
-  EAtIndex _ v _ix -> do
-    v' <- memoizedCompile v
-    pure . V.singleton $ v' V.! (fromIntegral _ix)
-  EUpdateIndex _ p b v -> do
-    v' <- memoizedCompile v
-    b' <- memoizedCompile b >>= assertSingleSource
-    let p' = fromIntegral p
-    pure $ V.imap (\_ix w -> if _ix == p' then b' else w) v'
 
 memoizedCompile ::
   forall f m.
@@ -385,32 +377,22 @@ exprToArithCircuit expr output = do
   compileOut <- memoizedCompile e >>= assertSingleSource
   emit $ Mul (ConstGate 1) (addVar compileOut) output
 
-fieldToBool 
-  :: (Hashable f, GaloisField f) =>  
-    TExpr.Expr Wire f f -> 
-    ExprM f (TExpr.Expr Wire f Bool)
+fieldToBool ::
+  (Hashable f, GaloisField f) =>
+  TExpr.Expr Wire f f ->
+  ExprM f (TExpr.Expr Wire f Bool)
 fieldToBool e = do
   eOut <- hashCons <$> unType e
   a <- memoizedCompile eOut >>= assertSingleSource >>= addWire
   emit $ Boolean a
   pure $ unsafeCoerce e
 
-unBundle ::
-  forall n f ty.
-  (KnownNat n, GaloisField f, Hashable f) =>
-  TExpr.Expr Wire f (SV.Vector n ty) ->
-  ExprM f (SV.Vector n (TExpr.Expr Wire f f))
-unBundle b = do
-  bis <- memoizedCompile . hashCons =<< unType b
-  ws <- traverse addWire bis
-  pure $ fromJust $ SV.toSized (TExpr.EVar . TExpr.VarField <$> ws)
-
-unType :: forall f ty m. MonadState (BuilderState f) m => TExpr.Expr Wire f ty -> m (Expr () Wire f)
+unType :: forall f ty m. (MonadState (BuilderState f) m) => TExpr.Expr Wire f ty -> m (Expr () Wire f)
 unType = \case
   TExpr.EVal v -> pure $ case v of
     TExpr.ValBool b -> EVal () b
     TExpr.ValField f -> EVal () f
-  TExpr.EVar v -> case v of 
+  TExpr.EVar v -> case v of
     TExpr.VarField w -> pure $ EVar () (UVar w)
     TExpr.VarBool b -> do
       emit $ Boolean b
@@ -421,8 +403,6 @@ unType = \case
   TExpr.EEq l r -> EEq () <$> unType l <*> unType r
   TExpr.ESplit i -> ESplit () (fromIntegral $ natVal (Proxy @(TExpr.NBits f))) <$> unType i
   TExpr.EJoin i -> EJoin () <$> unType i
-  TExpr.EAtIndex v ix -> EAtIndex () <$> unType v <*> pure (fromIntegral ix)
-  TExpr.EUpdateIndex p b v -> EUpdateIndex () (fromIntegral p) <$> unType b <*> unType v
   TExpr.EBundle b -> EBundle () <$> traverse unType (SV.fromSized b)
   where
     untypeBinOp :: TExpr.BinOp f a -> BinOp
@@ -440,3 +420,14 @@ unType = \case
       TExpr.UNeg -> UNeg
       TExpr.UNot -> UNot
 
+_unBundle ::
+  forall n f ty.
+  (KnownNat n) =>
+  (GaloisField f) =>
+  (Hashable f) =>
+  TExpr.Expr Wire f (SV.Vector n ty) ->
+  ExprM f (SV.Vector n (TExpr.Expr Wire f f))
+_unBundle b = do
+  bis <- memoizedCompile . hashCons =<< unType b
+  ws <- traverse addWire bis
+  pure $ fromJust $ SV.toSized (TExpr.EVar . TExpr.VarField <$> ws)

language/src/Circuit/Language/Compile.hs

@@ -1,9 +1,10 @@
-{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeFamilyDependencies #-}
 
 -- | Surface language
 module Circuit.Language.DSL
   ( Signal,
-    Bundle,
+    Bundled (..),
+    type Unbundled,
     cField,
     cBool,
     add,
@@ -25,7 +26,6 @@ module Circuit.Language.DSL
     joinBits,
     atIndex,
     updateIndex_,
-    bundle,
 
     -- * Monoids
     Any_ (..),
@@ -42,14 +42,14 @@ import Circuit.Language.Compile
 import Circuit.Language.TExpr
 import Data.Field.Galois (GaloisField, PrimeField)
 import Data.Finite (Finite)
-import Data.Vector.Sized (Vector)
+import Data.Vector.Sized (Vector, ix)
+import Lens.Micro ((.~), (^.))
 import Protolude
+import Unsafe.Coerce (unsafeCoerce)
 
 --------------------------------------------------------------------------------
 type Signal f = Expr Wire f
 
-type Bundle f n a = Expr Wire f (Vector n a)
-
 -- | Convert constant to expression
 cField :: f -> Signal f f
 cField = EVal . ValField
@@ -96,10 +96,10 @@ cond = EIf
 splitBits ::
   (KnownNat (NBits f)) =>
   Signal f f ->
-  Bundle f (NBits f) Bool
+  Signal f (Vector (NBits f) Bool)
 splitBits = ESplit
 
-joinBits :: (KnownNat n) => Bundle f n Bool -> Signal f f
+joinBits :: (KnownNat n) => Signal f (Vector n Bool) -> Signal f f
 joinBits = EJoin
 
 deref :: Var Wire f ty -> Signal f ty
@@ -111,14 +111,25 @@ retBool label sig = compileWithWire (boolInput Public label) sig
 retField :: (PrimeField f, Hashable f) => Text -> Signal f f -> ExprM f (Var Wire f f)
 retField label sig = compileWithWire (fieldInput Public label) sig
 
-atIndex :: (KnownNat n) => Bundle f n ty -> Finite n -> Signal f ty
-atIndex = EAtIndex
-
-updateIndex_ :: (KnownNat n) => Finite n -> Signal f ty -> Bundle f n ty -> Bundle f n ty
-updateIndex_ p = EUpdateIndex p
-
-bundle :: Vector n (Signal f ty) -> Bundle f n ty
-bundle = EBundle
+atIndex ::
+  (Bundled f (Vector n ty)) =>
+  Finite n ->
+  Signal f (Vector n ty) ->
+  ExprM f (Signal f ty)
+atIndex i b = do
+  bs <- unbundle b
+  return $ bs ^. ix i
+
+updateIndex_ ::
+  (Bundled f (Vector n ty)) =>
+  Finite n ->
+  Signal f ty ->
+  Signal f (Vector n ty) ->
+  ExprM f (Signal f (Vector n ty))
+updateIndex_ p s v = do
+  bs <- unbundle v
+  let bs' = bs & ix p .~ s
+  return $ bundle bs'
 
 --------------------------------------------------------------------------------
 
@@ -140,10 +151,10 @@ instance (Num f) => Monoid (Any_ f) where
 
 newtype Add_ f = Add_ {unAdd_ :: Signal f f}
 
-instance GaloisField f => Semigroup (Add_ f) where
-  Add_ a <> Add_ b = Add_ $ add a b 
+instance (GaloisField f) => Semigroup (Add_ f) where
+  Add_ a <> Add_ b = Add_ $ add a b
 
-instance GaloisField f => Monoid (Add_ f) where
+instance (GaloisField f) => Monoid (Add_ f) where
   mempty = Add_ $ cField 0
 
 --------------------------------------------------------------------------------
@@ -170,3 +181,20 @@ any_ ::
   t a ->
   Signal f Bool
 any_ f = unAny_ . foldMap (Any_ . f)
+
+--------------------------------------------------------------------------------
+
+type family Unbundled f a = res | res -> a where
+  Unbundled f (Vector n ty) = Vector n (Signal f ty)
+
+class Bundled f a where
+  bundle :: Unbundled f a -> Signal f a
+  unbundle :: Signal f a -> ExprM f (Unbundled f a)
+
+instance (Hashable f, GaloisField f, KnownNat n) => Bundled f (Vector n f) where
+  bundle = EBundle
+  unbundle = _unBundle
+
+instance (Hashable f, GaloisField f, KnownNat n) => Bundled f (Vector n Bool) where
+  bundle = EBundle
+  unbundle = fmap unsafeCoerce . _unBundle