tighten the belt

Author: martyall

Diff for: arithmetic-circuits.cabal

@@ -36,7 +36,7 @@ common extensions
     RecordWildCards
 
 common warnings
-  ghc-options: -Wall -Wredundant-constraints
+  ghc-options: -Wall -Wredundant-constraints -Werror
 
 
 common deps
@@ -87,7 +87,6 @@ library language
       arithmetic-circuits
     , array
     , hashable
-    , intern
     , microlens
     , mtl
     , unordered-containers
@@ -154,6 +153,7 @@ test-suite language-tests
     , crypton
     , bytestring
     , memory
+    , unordered-containers
 
 
   build-tool-depends: tasty-discover:tasty-discover >=4.2
@@ -169,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
@@ -235,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

@@ -1,7 +1,9 @@
 tests: True
+benchmarks: True
 
 packages: .
 
+
 source-repository-package
   type: git
   location: https://github.com/chessai/semirings.git
@@ -19,9 +21,3 @@ source-repository-package
   location: https://github.com/martyall/propagators.git
   tag: 6c1778171ee7c37ca29cf1ec4efaf79c5ec8af62
   --sha256: 6MUrOeUvp/iRFCTdvAyAGT2ymm5yKRqfN4GCkZ4eDyo=
-
-source-repository-package
-  type: git
-  location: https://github.com/l-adic/intern.git
-  tag: 5d0df17db13978d16533530370c2f7fd51e37cc8
-  --sha256: 6MUrOeUvp/iRFCTdvAyAGT2ymm5yKRqfN4GCkZ4eDyo=

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"
@@ -28,7 +28,7 @@ type Fr = Prime 2188824287183927522224640574525727508854836440041603434369820418
 
 program :: ExprM Fr (Var Wire Fr Bool)
 program = do
-  n <- deref <$> fieldInput Public "n"
-  a <- deref <$> fieldInput Private "a"
-  b <- deref <$> fieldInput Private "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 #-}

Diff for: circom-compat/src/R1CS/Circom.hs

@@ -28,7 +28,6 @@ data AffineCircuit f i
   | ScalarMul f (AffineCircuit f i)
   | ConstGate f
   | Var i
-  | Nil
   deriving (Read, Eq, Ord, Show, Generic, NFData)
 
 instance (FromJSON i, FromJSON f) => FromJSON (AffineCircuit f i)
@@ -47,16 +46,13 @@ 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 = prettyPrec 0
     where
       prettyPrec :: Int -> AffineCircuit f i -> Doc
       prettyPrec p e =
         case e of
-          Nil ->
-            text "nil"
           Var v ->
             pretty v
           ConstGate f -> pretty f
@@ -71,11 +67,13 @@ instance (Pretty i, Pretty f) => Pretty (AffineCircuit f i) where
 parensPrec :: Int -> Int -> Doc -> Doc
 parensPrec opPrec p = if p > opPrec then parens else identity
 
+{-# SCC evalAffineCircuit #-}
+
 -- | Evaluate the arithmetic circuit without mul-gates on the given
 -- input. Variable map is assumed to have all the variables referred
 -- to in the circuit. Failed lookups are currently treated as 0.
 evalAffineCircuit ::
-  (Num f) =>
+  (Num f, Show i) =>
   -- | lookup function for variable mapping
   (i -> vars -> Maybe f) ->
   -- | variables
@@ -84,9 +82,8 @@ evalAffineCircuit ::
   AffineCircuit f i ->
   f
 evalAffineCircuit lookupVar vars = \case
-  Nil -> 0
   ConstGate f -> f
-  Var i -> fromMaybe 0 $ lookupVar i vars
+  Var i -> fromMaybe (panic $ "missing variable assignment: " <> show i) $ lookupVar i vars
   Add l r -> evalAffineCircuit lookupVar vars l + evalAffineCircuit lookupVar vars r
   ScalarMul scalar expr -> evalAffineCircuit lookupVar vars expr * scalar
 
@@ -99,7 +96,6 @@ affineCircuitToAffineMap ::
   -- | constant part and non-constant part
   (f, Map i f)
 affineCircuitToAffineMap = \case
-  Nil -> (0, mempty)
   Var i -> (0, Map.singleton i 1)
   Add l r -> (constLeft + constRight, Map.unionWith (+) vecLeft vecRight)
     where