Merge pull request #7 from l-adic/middle-out

Middle out

Author: martyall

Diff for: arithmetic-circuits.cabal

@@ -79,19 +79,21 @@ library language
   exposed-modules:
     Circuit.Language.Compile
     Circuit.Language.Expr
-    Circuit.Language.TExpr
     Circuit.Language.DSL
     Circuit.Language
 
   hs-source-dirs:     language/src
   build-depends:
       arithmetic-circuits
+    , array
     , hashable
     , microlens
+    , mtl
+    , unordered-containers
     , vector
     , vector-sized
 
-  ghc-options:        -freverse-errors -O2 -Wall
+  ghc-options:        -freverse-errors  -O2 -Wall
 
   default-language:   GHC2021
 
@@ -125,16 +127,18 @@ test-suite language-tests
     Test.Circuit.Expr
     Test.Circuit.Sudoku
     Test.Circuit.Lang
+    Test.Circuit.SHA3
 
   hs-source-dirs:     language/test
 
-  ghc-options:        -freverse-errors -O2 -Wall -main-is Main
+  ghc-options:        -freverse-errors  -O2 -Wall -main-is Main
   build-depends:
       arithmetic-circuits
     , arithmetic-circuits:language
     , array
     , distributive 
     , fin
+    , microlens
     , quickcheck-instances  >=0.3
     , QuickCheck
     , random
@@ -144,6 +148,13 @@ test-suite language-tests
     , tasty-hspec
     , tasty-quickcheck      >=0.10
     , vec
+    , vector
+    , vector-sized
+    , crypton
+    , bytestring
+    , memory
+    , unordered-containers
+
 
   build-tool-depends: tasty-discover:tasty-discover >=4.2
 
@@ -158,9 +169,11 @@ executable factors
     , arithmetic-circuits
     , arithmetic-circuits:language
     , arithmetic-circuits:circom-compat
+    , vector
 
   hs-source-dirs:   circom-compat/app
   default-language: GHC2021
+  ghc-options:      -freverse-errors -O2 -Wall -main-is Main
 
 test-suite circuit-tests
   import: deps, extensions
@@ -224,14 +237,16 @@ benchmark circuit-benchmarks
   type:               exitcode-stdio-1.0
   main-is:            Main.hs
   other-modules:
-    Circuit
+    Circuit.Bench
     Paths_arithmetic_circuits
 
   hs-source-dirs:     bench
 
   ghc-options:        -freverse-errors -O2 -Wall -main-is Main
   build-depends:
       arithmetic-circuits
-    , criterion            >=1.6
+    , arithmetic-circuits:language
+    , criterion
+    , vector
 
   default-language:   GHC2021

Diff for: bench/Circuit.hs

@@ -0,0 +1,54 @@
+{-# LANGUAGE DataKinds #-}
+
+module Circuit.Bench where
+
+import Circuit
+import Circuit.Language
+import Criterion
+import Data.Field.Galois (Prime)
+import Data.IntMap qualified as IntMap
+import Data.Map qualified as Map
+import Data.Vector (generateM)
+import GHC.TypeNats (Natural, SNat, withKnownNat, withSomeSNat)
+import Protolude
+
+type Fr = Prime 21888242871839275222246405745257275088548364400416034343698204186575808495617
+
+benchmarks :: Benchmark
+benchmarks =
+  bgroup
+    "largeMult"
+    [ bench "1_000" $ whnf largeMult 1000,
+      bench "10_000" $ whnf largeMult 10000,
+      bench "100_000" $ whnf largeMult 100_000,
+      bench "1_000_000" $ whnf largeMult 1_000_000
+    ]
+
+largeMult :: Natural -> Fr
+largeMult n =
+  withSomeSNat n $ \(sn :: SNat n) ->
+    withKnownNat sn $
+      let BuilderState {bsVars, bsCircuit} = snd $ runCircuitBuilder (program (Proxy @n))
+          inputs =
+            assignInputs bsVars $
+              Map.fromList $
+                map (\i -> (nameInput i, fromIntegral i + 1)) [0 .. n - 1]
+          w = altSolve bsCircuit inputs
+       in fromMaybe (panic "output not found") $ lookupVar bsVars "out" w
+
+nameInput :: (Integral a) => a -> Text
+nameInput i = "x" <> show (toInteger i)
+
+program :: forall n. (KnownNat n) => Proxy n -> ExprM Fr (Var Wire Fr Fr)
+program p = do
+  xs <- generateM (fromIntegral $ natVal p) $ \i ->
+    var_ <$> fieldInput Public (nameInput i)
+  fieldOutput "out" $ product xs
+
+altSolve :: ArithCircuit Fr -> IntMap Fr -> IntMap Fr
+altSolve p inputs =
+  evalArithCircuit
+    (\w m -> IntMap.lookup (wireName w) m)
+    (\w m -> IntMap.insert (wireName w) m)
+    p
+    inputs

Diff for: bench/Circuit/Bench.hs

@@ -4,12 +4,12 @@
 
 module Main where
 
-import Circuit qualified
+import Circuit.Bench qualified
 import Criterion.Main
 import Protolude
 
 main :: IO ()
 main =
   defaultMain
-    [ bgroup "Circuit to QAP translation" Circuit.benchmarks
+    [ Circuit.Bench.benchmarks
     ]

Diff for: bench/Main.hs

@@ -3,6 +3,7 @@ benchmarks: True
 
 packages: .
 
+
 source-repository-package
   type: git
   location: https://github.com/chessai/semirings.git
@@ -12,7 +13,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

Diff for: cabal.project

@@ -7,17 +7,17 @@ import Circuit
 import Circuit.Language
 import Data.Binary (encodeFile)
 import Data.Field.Galois (Prime)
-import Data.Map qualified as Map
-import Data.Set qualified as Set
+import Data.IntMap qualified as IntMap
+import Data.IntSet qualified as IntSet
 import Protolude
 import R1CS (Inputs (..), calculateWitness, isValidWitness)
 import R1CS.Circom (r1csToCircomR1CS, witnessToCircomWitness)
 
 main :: IO ()
 main = do
   let BuilderState {..} = snd $ runCircuitBuilder program
-      publicInputs = Map.fromList $ zip (Set.toAscList $ cvPublicInputs bsVars) [6]
-      privateInputs = Map.fromList $ zip (Set.toAscList $ cvPrivateInputs bsVars) [2, 3]
+      publicInputs = IntMap.fromList $ zip (IntSet.toAscList $ cvPublicInputs bsVars) [6]
+      privateInputs = IntMap.fromList $ zip (IntSet.toAscList $ cvPrivateInputs bsVars) [2, 3]
       inputs = publicInputs <> privateInputs
       (r1cs, wtns) = calculateWitness bsVars bsCircuit (Inputs inputs)
   unless (isValidWitness wtns r1cs) $ panic "Invalid witness"
@@ -26,9 +26,9 @@ main = do
 
 type Fr = Prime 21888242871839275222246405745257275088548364400416034343698204186575808495617
 
-program :: ExprM Fr Wire
+program :: ExprM Fr (Var Wire Fr Bool)
 program = do
-  n <- deref <$> fieldInput Public "n"
-  a <- deref <$> fieldInput Private "a"
-  b <- deref <$> fieldInput Private "b"
-  retBool "out" $ eq n (a * b)
+  n <- var_ <$> fieldInput Public "n"
+  a <- var_ <$> fieldInput Private "a"
+  b <- var_ <$> fieldInput Private "b"
+  boolOutput "out" $ eq_ n (a * b)

Diff for: circom-compat/app/Main.hs

@@ -21,8 +21,9 @@ where
 import Circuit
 import Data.Field.Galois (GaloisField, PrimeField (fromP), char)
 import Data.IORef (IORef, readIORef, writeIORef)
+import Data.IntMap qualified as IntMap
+import Data.IntSet qualified as IntSet
 import Data.Map qualified as Map
-import Data.Set qualified as Set
 import Data.Vector qualified as V
 import Data.Vector.Mutable (IOVector)
 import Data.Vector.Mutable qualified as MV
@@ -52,8 +53,8 @@ mkProgramEnv vars circ =
     { peFieldSize = FieldSize 32,
       peRawPrime = toInteger $ char (1 :: f),
       peVersion = 2,
-      peInputsSize = Set.size $ cvPrivateInputs vars <> cvPublicInputs vars,
-      peWitnessSize = Set.size $ Set.insert oneVar $ cvVars vars,
+      peInputsSize = IntSet.size $ cvPrivateInputs vars <> cvPublicInputs vars,
+      peWitnessSize = IntSet.size $ IntSet.insert oneVar $ cvVars vars,
       peCircuit = circ,
       peCircuitVars = relabel hashText vars
     }
@@ -125,16 +126,16 @@ _setInputSignal env@(ProgramEnv {peCircuit, peInputsSize, peCircuitVars}) stRef
   st <- readIORef stRef
   let Inputs inputs = psInputs st
   let h = mkFNV msb lsb
-      i = fromMaybe (panic $ "Hash not found: " <> show h) $ Map.lookup h (cvInputsLabels peCircuitVars)
+      i = fromMaybe (panic $ "Hash not found: " <> show h) $ Map.lookup h (labelToVar $ cvInputsLabels peCircuitVars)
   newInput <- fromInteger <$> readBuffer env stRef
-  let newInputs = Map.insert i newInput inputs
+  let newInputs = IntMap.insert i newInput inputs
   writeIORef stRef $
-    if Map.size newInputs == peInputsSize
+    if IntMap.size newInputs == peInputsSize
       then
         let wtns = solve peCircuitVars peCircuit newInputs
          in st
               { psInputs = Inputs newInputs,
-                psWitness = Witness $ Map.insert oneVar 1 wtns
+                psWitness = Witness $ IntMap.insert oneVar 1 wtns
               }
       else st {psInputs = Inputs newInputs}
 
@@ -149,7 +150,7 @@ _getWitness ::
   IO ()
 _getWitness env stRef i = do
   ProgramState {psWitness = Witness wtns} <- readIORef stRef
-  let wtn = maybe (panic $ "missing witness " <> show i) fromP $ Map.lookup i wtns
+  let wtn = maybe (panic $ "missing witness " <> show i) fromP $ IntMap.lookup i wtns
    in writeBuffer env stRef wtn
 
 --------------------------------------------------------------------------------

Diff for: circom-compat/src/Circuit/Solver/Circom.hs

@@ -21,7 +21,7 @@ import Data.Binary.Get (getInt32le, getInt64le, getWord32le, getWord64le, lookAh
 import Data.Binary.Put (putInt32le, putLazyByteString, putWord32le, putWord64le, runPut)
 import Data.ByteString.Lazy qualified as LBS
 import Data.Field.Galois (GaloisField (char), PrimeField, fromP)
-import Data.Map qualified as Map
+import Data.IntMap qualified as IntMap
 import Data.Vector (Vector)
 import Data.Vector qualified as V
 import Protolude
@@ -261,14 +261,14 @@ getPoly fieldSize = do
   LinearCombination factors <- getLinearCombination fieldSize
   pure $
     LinearPoly $
-      foldl (\acc (Factor {wireId, value}) -> Map.insert (fromIntegral wireId) value acc) mempty factors
+      foldl (\acc (Factor {wireId, value}) -> IntMap.insert (fromIntegral wireId) value acc) mempty factors
 
 putPoly :: (PrimeField k) => FieldSize -> LinearPoly k -> Put
 putPoly fieldSize (LinearPoly p) =
   putLinearCombination fieldSize $
     LinearCombination
       [ Factor {wireId = fromIntegral var, value}
-        | (var, value) <- Map.toAscList p,
+        | (var, value) <- IntMap.toAscList p,
           value /= 0
       ]
 
@@ -325,14 +325,14 @@ witnessToCircomWitness (Witness m) =
         WitnessHeader
           { whFieldSize = FieldSize 32,
             whPrime = fromIntegral $ char (1 :: f),
-            whWitnessSize = fromIntegral $ Map.size m
+            whWitnessSize = fromIntegral $ IntMap.size m
           },
-      wtnsValues = snd <$> Map.toAscList m
+      wtnsValues = snd <$> IntMap.toAscList m
     }
 
 witnessFromCircomWitness :: CircomWitness f -> Witness f
 witnessFromCircomWitness (CircomWitness {wtnsValues}) =
-  Witness $ Map.fromList $ zip [0 ..] wtnsValues
+  Witness $ IntMap.fromList $ zip [0 ..] wtnsValues
 
 instance (PrimeField k) => Binary (CircomWitness k) where
   get = do
@@ -425,6 +425,7 @@ putWitnessValues fieldSize values = do
 integerFromLittleEndian :: Vector Word32 -> Integer
 integerFromLittleEndian bytes =
   foldl' (\acc (i, byte) -> acc .|. (fromIntegral byte `shiftL` (i * 32))) 0 (V.zip (V.fromList [0 ..]) bytes)
+{-# INLINE integerFromLittleEndian #-}
 
 integerToLittleEndian :: FieldSize -> Integer -> Vector Word32
 integerToLittleEndian fieldSize n =
@@ -434,3 +435,4 @@ integerToLittleEndian fieldSize n =
   where
     go 0 = mempty
     go x = fromIntegral (x .&. 0xffffffff) `V.cons` go (x `shiftR` 32)
+{-# INLINE integerToLittleEndian #-}