proper signal size

Author: martyall

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

@@ -2,7 +2,7 @@ module Circom.CLI (defaultMain) where
 
 import Circom.R1CS (R1CSHeader (rhPrime), decodeR1CSHeaderFromFile, r1csFromCircomR1CS, r1csToCircomR1CS, witnessFromCircomWitness)
 import Circom.Solver (CircomProgram (..), mkCircomProgram, nativeGenWitness)
-import Circuit.Arithmetic (CircuitVars (..), InputBindings (labelToVar), VarType (..), restrictVars)
+import Circuit.Arithmetic (CircuitVars (..), InputBindings (labelToVar), VarType (..), inputSizes, restrictVars)
 import Circuit.Dataflow qualified as DataFlow
 import Circuit.Dot (arithCircuitToDot)
 import Circuit.Language.Compile (BuilderState (..), ExprM, runCircuitBuilder)
@@ -21,7 +21,6 @@ import GHC.TypeNats (SNat, withKnownNat, withSomeSNat)
 import Numeric (showHex)
 import Options.Applicative (CommandFields, Mod, Parser, ParserInfo, command, eitherReader, execParser, fullDesc, header, help, helper, hsubparser, info, long, option, progDesc, showDefault, showDefaultWith, strOption, switch, value)
 import Protolude
-import Protolude.Unsafe (unsafeHead)
 import R1CS (R1CS, Witness (Witness), isValidWitness, toR1CS)
 import Prelude (MonadFail (fail))
 
@@ -328,11 +327,7 @@ decodeIOVars enc v = do
 mkInputsTemplate :: Encoding -> CircuitVars Text -> IOVars
 mkInputsTemplate enc vars =
   let inputsOnly = cvInputsLabels $ restrictVars vars (cvPrivateInputs vars `IntSet.union` cvPublicInputs vars)
-      vs =
-        map (\a -> (fst $ unsafeHead a, length a)) $
-          groupBy (\a b -> fst a == fst b) $
-            Map.keys $
-              labelToVar inputsOnly
+      vs = Map.toList $ inputSizes inputsOnly
       f (label, len) =
         if len > 1
           then (label, Array (replicate len 0))

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

@@ -77,6 +77,7 @@ data ProgramEnv f = ProgramEnv
     peInputsSize :: Int,
     peWitnessSize :: Int,
     peCircuit :: ArithCircuit f,
+    peSignalSizes :: Map FNVHash Int,
     peCircuitVars :: CircuitVars FNVHash
   }
 
@@ -86,14 +87,16 @@ mkProgramEnv ::
   CircomProgram f ->
   ProgramEnv f
 mkProgramEnv CircomProgram {cpVars = vars, cpCircuit = circ} =
-  ProgramEnv
+  let vs = relabel hashText vars
+  in ProgramEnv
     { peFieldSize = FieldSize 32,
       peRawPrime = toInteger $ char (1 :: f),
       peVersion = 2,
       peInputsSize = IntSet.size $ cvPrivateInputs vars <> cvPublicInputs vars,
       peWitnessSize = IntSet.size $ IntSet.insert oneVar $ cvVars vars,
       peCircuit = circ,
-      peCircuitVars = relabel hashText vars
+      peSignalSizes = inputSizes (cvInputsLabels vs),
+      peCircuitVars = vs
     }
 
 data ProgramState f = ProgramState
@@ -150,8 +153,11 @@ _getInputSize :: ProgramEnv f -> Int
 _getInputSize = peInputsSize
 
 -- we dont (yet) support multiple values (e.g. arrays) for signal values
-_getInputSignalSize :: Word32 -> Word32 -> IO Int
-_getInputSignalSize _ _ = pure 1
+_getInputSignalSize :: ProgramEnv f -> Word32 -> Word32 -> IO Int
+_getInputSignalSize ProgramEnv {peSignalSizes} msb lsb = 
+  let h = mkFNV msb lsb
+  in pure $ fromMaybe 0 $ Map.lookup h peSignalSizes
+  
 
 -- we ignore the last arugment because our signals don't have indices, only names
 _setInputSignal ::

Diff for: circuit/src/Circuit/Arithmetic.hs

@@ -23,6 +23,7 @@ module Circuit.Arithmetic
     nGates,
     InputBindings (..),
     insertInputBinding,
+    inputSizes,
     Reindexable (..),
     restrictVars,
   )
@@ -50,6 +51,7 @@ import Text.PrettyPrint.Leijen.Text as PP
     vcat,
     (<+>),
   )
+import Protolude.Unsafe (unsafeHead)
 
 data InputType = Public | Private deriving (Show, Eq, Ord, Generic, NFData)
 
@@ -489,6 +491,13 @@ restrictInputBindings s InputBindings {..} =
       varToLabel = IntMap.restrictKeys varToLabel s
     }
 
+inputSizes :: (Ord label) => InputBindings label -> Map label Int
+inputSizes InputBindings {..} =
+  Map.fromList $
+    map (\a -> (fst $ unsafeHead a, length a)) $
+      groupBy (\a b -> fst a == fst b) $
+        Map.keys labelToVar
+
 --------------------------------------------------------------------------------
 
 class Reindexable a where