added vector ops for bitvec operations

martyall · martyall · commit 7785d46604b5 · 2024-05-07T09:38:49.000-07:00
diff --git a/language/src/Circuit/Language/DSL.hs b/language/src/Circuit/Language/DSL.hs
@@ -11,9 +11,13 @@ module Circuit.Language.DSL
     sub,
     mul,
     and_,
+    ands_,
     or_,
+    ors_,
     xor_,
+    xors_,
     not_,
+    nots_,
     eq_,
     fieldInput,
     boolInput,
@@ -80,10 +84,22 @@ and_ = binOp_ BAnd
 or_ = binOp_ BOr
 xor_ = binOp_ BXor
 
+ands_, ors_, xors_ ::
+  (Hashable f) =>
+  Signal f ('TVec n 'TBool) ->
+  Signal f ('TVec n 'TBool) ->
+  Signal f ('TVec n 'TBool)
+ands_ = binOp_ BAnds
+ors_ = binOp_ BOrs
+xors_ = binOp_ BXors
+
 -- | Negate expression
 not_ :: (Hashable f) => Signal f 'TBool -> Signal f 'TBool
 not_ = unOp_ UNot
 
+nots_ :: (Hashable f) => Signal f ('TVec n 'TBool) -> Signal f ('TVec n 'TBool)
+nots_ = unOp_ UNots
+
 fieldInput :: InputType -> Text -> ExprM f (Var Wire f 'TField)
 fieldInput it label =
   case it of
@@ -106,10 +122,10 @@ fieldOutput label s = do
 fieldsOutput :: (KnownNat n, Hashable f, GaloisField f) => Vector n (Var Wire f 'TField) -> Signal f ('TVec n 'TField) -> ExprM f (Vector n (Var Wire f 'TField))
 fieldsOutput vs s = fromJust . SV.toSized <$> compileWithWires (SV.fromSized vs) s
 
-boolOutput :: (Hashable f, GaloisField f) => Text -> Signal f 'TBool -> ExprM f (Var Wire f 'TBool)
+boolOutput :: forall f. (Hashable f, GaloisField f) => Text -> Signal f 'TBool -> ExprM f (Var Wire f 'TBool)
 boolOutput label s = do
   out <- VarBool <$> freshOutput label
-  unsafeCoerce <$> compileWithWire (boolToField out) (boolToField s)
+  unsafeCoerce <$> compileWithWire (boolToField @(Var Wire f 'TBool) out) (boolToField s)
 {-# INLINE boolOutput #-}
 
 boolsOutput :: (KnownNat n, Hashable f, GaloisField f) => Vector n (Var Wire f 'TBool) -> Signal f ('TVec n 'TBool) -> ExprM f (Vector n (Var Wire f 'TBool))
@@ -213,7 +229,7 @@ any_ f = unAny_ . foldMap (Any_ . f)
 --------------------------------------------------------------------------------
 
 class Bundle f a where
-  type Unbundled f a
+  type Unbundled f a = r | r -> a
   bundle :: Unbundled f a -> Signal f a
   unbundle :: Signal f a -> ExprM f (Unbundled f a)
 
diff --git a/language/src/Circuit/Language/Expr.hs b/language/src/Circuit/Language/Expr.hs
@@ -35,7 +35,7 @@ module Circuit.Language.Expr
   )
 where
 
-import Data.Field.Galois (Prime, PrimeField (fromP))
+import Data.Field.Galois (Prime, PrimeField (fromP), GaloisField)
 import Data.Semiring (Ring (..), Semiring (..))
 import Data.Sequence ((|>))
 import Data.Set qualified as Set
@@ -91,6 +91,7 @@ rawWire (VarBool i) = i
 data UnOp f (ty :: Ty) where
   UNeg :: UnOp f 'TField
   UNot :: UnOp f 'TBool
+  UNots :: UnOp f ('TVec n 'TBool)
 
 deriving instance (Show f) => Show (UnOp f a)
 
@@ -100,15 +101,19 @@ instance Pretty (UnOp f a) where
   pretty op = case op of
     UNeg -> text "neg"
     UNot -> text "!"
+    UNots -> text "nots"
 
 data BinOp f (a :: Ty) where
   BAdd :: BinOp f 'TField
   BSub :: BinOp f 'TField
   BMul :: BinOp f 'TField
   BDiv :: BinOp f 'TField
   BAnd :: BinOp f 'TBool
+  BAnds :: BinOp f (TVec n 'TBool)
   BOr :: BinOp f 'TBool
+  BOrs :: BinOp f (TVec n 'TBool)
   BXor :: BinOp f 'TBool
+  BXors :: BinOp f (TVec n 'TBool)
 
 deriving instance (Show f) => Show (BinOp f a)
 
@@ -121,13 +126,19 @@ instance Pretty (BinOp f a) where
     BMul -> text "*"
     BDiv -> text "/"
     BAnd -> text "&&"
+    BAnds -> text ".&&."
     BOr -> text "||"
+    BOrs -> text ".||."
     BXor -> text "xor"
+    BXors -> text "xors"
 
 opPrecedence :: BinOp f a -> Int
 opPrecedence BOr = 5
+opPrecedence BOrs = 5
 opPrecedence BXor = 5
+opPrecedence BXors = 5
 opPrecedence BAnd = 5
+opPrecedence BAnds = 5
 opPrecedence BSub = 6
 opPrecedence BAdd = 6
 opPrecedence BMul = 7
@@ -246,6 +257,8 @@ evalExpr lookupVar vars expr = case expr of
     Protolude.negate $ evalExpr lookupVar vars e1
   EUnOp _ UNot e1 ->
     not $ evalExpr lookupVar vars e1
+  EUnOp _ UNots e1 ->
+    SV.map not $ evalExpr lookupVar vars e1
   EBinOp _ op e1 e2 ->
     let e1' = evalExpr lookupVar vars e1
         e2' = evalExpr lookupVar vars e2
@@ -257,8 +270,11 @@ evalExpr lookupVar vars expr = case expr of
         BMul -> (*)
         BDiv -> (/)
         BAnd -> (&&)
+        BAnds -> SV.zipWith (&&)
         BOr -> (||)
+        BOrs -> SV.zipWith (||)
         BXor -> \x y -> (x || y) && not (x && y)
+        BXors -> SV.zipWith (\x y -> (x || y) && not (x && y))
   EIf _ b true false ->
     let cond = evalExpr lookupVar vars b
      in if cond
@@ -401,9 +417,12 @@ bundle_ b =
    in EBundle h b
 {-# INLINE bundle_ #-}
 
-class BoolToField b f | b -> f where
+class BoolToField b f where
   boolToField :: b -> f
 
+instance GaloisField f => BoolToField Bool f where
+  boolToField b = fromInteger $ if b then 1 else 0
+
 instance BoolToField (Val f 'TBool) (Val f 'TField) where
   boolToField (ValBool b) = ValField b
 
@@ -418,7 +437,14 @@ instance BoolToField (Expr i f ('TVec n 'TBool)) (Expr i f ('TVec n 'TField)) wh
 
 -------------------------------------------------------------------------------
 
-data UBinOp = UBAdd | UBSub | UBMul | UBDiv | UBAnd | UBOr | UBXor deriving (Show, Eq, Generic)
+data UBinOp = 
+  UBAdd | 
+  UBSub | 
+  UBMul | 
+  UBDiv | 
+  UBAnd | 
+  UBOr | 
+  UBXor deriving (Show, Eq, Generic)
 
 instance Hashable UBinOp
 
@@ -481,6 +507,7 @@ instance (Hashable i, Hashable f) => Hashable (Node i f) where
 untypeUnOp :: UnOp f a -> UUnOp
 untypeUnOp UNeg = UUNeg
 untypeUnOp UNot = UUNot
+untypeUnOp UNots = UUNot
 {-# INLINE untypeUnOp #-}
 
 untypeBinOp :: BinOp f a -> UBinOp
@@ -489,8 +516,11 @@ untypeBinOp BSub = UBSub
 untypeBinOp BMul = UBMul
 untypeBinOp BDiv = UBDiv
 untypeBinOp BAnd = UBAnd
+untypeBinOp BAnds = UBAnd
 untypeBinOp BOr = UBOr
+untypeBinOp BOrs = UBOr
 untypeBinOp BXor = UBXor
+untypeBinOp BXors = UBXor
 {-# INLINE untypeBinOp #-}
 
 --------------------------------------------------------------------------------
diff --git a/language/test/Test/Circuit/Lang.hs b/language/test/Test/Circuit/Lang.hs
@@ -11,7 +11,8 @@ import Data.Finite (Finite)
 import Data.IntMap qualified as IntMap
 import Data.Map qualified as Map
 import Protolude
-import Test.QuickCheck (Property, (.&&.), (=/=), (===), (==>))
+import Data.Vector.Sized qualified as SV
+import Test.QuickCheck (Arbitrary (arbitrary), Property, forAll, vectorOf, (.&&.), (=/=), (===), (==>))
 
 type Fr = Prime 21888242871839275222246405745257275088548364400416034343698204186575808495617
 
@@ -142,4 +143,78 @@ prop_sharingProg x y =
       computed = evalExpr IntMap.lookup input (relabelExpr wireName prog)
    in res === Just expected .&&. computed === expected
 
+binopsProg ::
+  BinOp Fr ('TVec 50 'TBool) ->
+  ExprM Fr (SV.Vector 50 (Var Wire Fr 'TBool))
+binopsProg op = do
+  xs <- SV.generateM $ \i ->
+    var_ <$> boolInput Public ("x" <> showFinite i)
+  ys <- SV.generateM $ \i ->
+    var_ <$> boolInput Public ("y" <> showFinite i)
+  zs <- SV.generateM $ \i ->
+    VarBool <$> freshOutput ("out" <> showFinite i)
+  boolsOutput zs $ binOp_ op (bundle xs) (bundle ys)
+
+propBinopsProg ::
+  BinOp Fr ('TVec 50 'TBool) ->
+  (Bool -> Bool -> Bool) ->
+  Property
+propBinopsProg top op = forAll arbInputs $ \(bs, bs') ->
+  let _xs = zip (map (\i -> "x" <> show @Int i) [0 ..]) bs
+      _ys = zip (map (\i -> "y" <> show @Int i) [0 ..]) bs'
+      (_, BuilderState {bsVars, bsCircuit}) = runCircuitBuilder (binopsProg top)
+      inputs =
+        assignInputs bsVars $
+          fmap boolToField $
+            Map.fromList (_xs <> _ys)
+      w = solve bsVars bsCircuit inputs
+      expected = map boolToField $ zipWith op bs bs'
+   in all (\(i, b) -> lookupVar bsVars ("out" <> show @Int i) w == Just b) $
+        zip [0 ..] expected
+  where
+    arbInputs = ((,) <$> vectorOf 50 arbitrary <*> vectorOf 50 arbitrary)
+
+prop_andsProg :: Property
+prop_andsProg = propBinopsProg BAnds (&&)
+
+prop_orsProg :: Property
+prop_orsProg = propBinopsProg BOrs (||)
+
+prop_xorsProg :: Property
+prop_xorsProg = propBinopsProg BXors (/=)
+
+unopsProg ::
+  UnOp Fr ('TVec 50 'TBool) ->
+  ExprM Fr (SV.Vector 50 (Var Wire Fr 'TBool))
+unopsProg op = do
+  xs <- SV.generateM $ \i ->
+    var_ <$> boolInput Public ("x" <> showFinite i)
+  zs <- SV.generateM $ \i ->
+    VarBool <$> freshOutput ("out" <> showFinite i)
+  boolsOutput zs $ unOp_ op (bundle xs)
+
+propUnopsProg ::
+  UnOp Fr ('TVec 50 'TBool) ->
+  (Bool -> Bool) ->
+  Property
+propUnopsProg top op = forAll arbInputs $ \bs ->
+  let _xs = zip (map (\i -> "x" <> show @Int i) [0 ..]) bs
+      (_, BuilderState {bsVars, bsCircuit}) = runCircuitBuilder (unopsProg top)
+      inputs =
+        assignInputs bsVars $
+          fmap boolToField $
+            Map.fromList _xs
+      w = solve bsVars bsCircuit inputs
+      expected = map (boolToField . op) bs
+   in all (\(i, b) -> lookupVar bsVars ("out" <> show @Int i) w == Just b) $
+        zip [0 ..] expected
+  where
+    arbInputs = vectorOf 50 arbitrary
+
+prop_notsProg :: Property
+prop_notsProg = propUnopsProg UNots not
+
 --------------------------------------------------------------------------------
+
+showFinite :: (KnownNat n) => Finite n -> Text
+showFinite = show . toInteger
diff --git a/language/test/Test/Circuit/SHA3.hs b/language/test/Test/Circuit/SHA3.hs
