Merge pull request #13 from torsion-labs/public-output

allow public assignment to output variables

.github/workflows/nix-ci.yml

@@ -18,7 +18,7 @@ jobs:
       uses: jaxxstorm/[email protected]
       with:
         repo: torsion-labs/arkworks-bridge
-        tag: v1.0.0-rc1
+        tag: v1.0.0-rc2
 
     - uses: cachix/install-nix-action@v24
       with:

app/Main.hs

@@ -10,4 +10,4 @@ import qualified Test.Snarkl.Unit.Programs as Programs
 import Prelude hiding (return, (+), (>>=))
 
 main :: IO ()
-main = defaultMain "prog" (Programs.prog2 1 :: Comp 'TField F_BN128)
+main = defaultMain "prog" Programs.prog1

src/Snarkl/CLI/Compile.hs

@@ -22,7 +22,7 @@ import Data.Typeable (Typeable)
 import Options.Applicative (CommandFields, Mod, Parser, command, execParser, fullDesc, header, help, helper, info, long, progDesc, showDefault, strOption, subparser, switch, value, (<**>))
 import Snarkl.AST (Comp, InputVariable (PrivateInput))
 import Snarkl.CLI.Common (mkConstraintsFilePath, mkInputsFilePath, mkR1CSFilePath, writeFileWithDir)
-import Snarkl.Common (InputVar (PrivateInputVar, PublicInputVar))
+import Snarkl.Common (InputVar (OutputVar, PrivateInputVar, PublicInputVar))
 import Snarkl.Compile
   ( SimplParam (RemoveUnreachable, Simplify),
     TExpPkg (TExpPkg),
@@ -100,9 +100,10 @@ compile CompileOpts {..} name comp = do
       (r1cs, scs, privateInputMap) = compileTExpToR1CS simpl texpPkg
       publicInputs = map PublicInputVar . cs_public_in_vars . unSimplifiedConstraintSystem $ scs
       privateInputs = map (uncurry PrivateInputVar) $ Map.toList privateInputMap
+      outputs = map OutputVar . cs_out_vars . unSimplifiedConstraintSystem $ scs
   let r1csFP = mkR1CSFilePath r1csOutput name
   writeFileWithDir r1csFP $ toJSONLines r1cs
   putStrLn $ "Wrote R1CS to " <> r1csFP
   let inputsFP = mkInputsFilePath inputsOutput name
-  writeFileWithDir inputsFP $ toJSONLines (publicInputs <> privateInputs)
+  writeFileWithDir inputsFP $ toJSONLines (publicInputs <> privateInputs <> outputs)
   putStrLn $ "Wrote inputs to " <> inputsFP

src/Snarkl/CLI/GenWitness.hs

@@ -82,13 +82,20 @@ genWitness GenWitnessOpts {..} name comp = do
   let (r1cs, constraints, _) = compileCompToR1CS simpl comp
   let [out_var] = cs_out_vars (unSimplifiedConstraintSystem constraints)
   -- parse the inputs, either from cli or from file
-  (pubInputs, privInputs) <- do
+  (pubInputs, privInputs, outputs) <- do
     let assignmentsFP = mkAssignmentsFilePath assignments name
     eInput <- fromJSONLines <$> readFileLines assignmentsFP
     let inputAssignments = either (failWith . ErrMsg) id eInput
     pure $ splitInputAssignments inputAssignments
-  let out_interp = comp_interp comp pubInputs (Map.mapKeys fst privInputs)
-  let witness@(Witness {witness_assgn = Assgn m}) = wit_of_cs pubInputs (Map.mapKeys snd privInputs) constraints
+  let mOut = case outputs of
+        [(v, k)] ->
+          if v /= out_var
+            then failWith $ ErrMsg $ "derived output variable " <> show out_var <> " does not match provided output variable " <> show v
+            else Just (v, k)
+        [] -> Nothing
+        _ -> failWith $ ErrMsg $ "expected one output variable, got " <> show outputs
+      initAssignments = Map.mapKeys snd privInputs <> maybe mempty (uncurry Map.singleton) mOut
+  let witness@(Witness {witness_assgn = Assgn m}) = wit_of_cs pubInputs initAssignments constraints
       out = case Map.lookup out_var m of
         Nothing ->
           failWith $
@@ -99,12 +106,13 @@ genWitness GenWitnessOpts {..} name comp = do
                   ++ show witness
               )
         Just out_val -> out_val
+      out_interp = comp_interp comp pubInputs (Map.mapKeys fst privInputs)
   unless (out_interp == out) $
     failWith $
       ErrMsg $
         "interpreter result "
           ++ show out_interp
-          ++ " differs from actual result "
+          ++ " differs from result computed from constraints"
           ++ show out
   let satisfies = sat_r1cs witness r1cs
   unless satisfies $

src/Snarkl/CLI/RunAll.hs

@@ -66,13 +66,20 @@ runAll RunAllOpts {..} name comp = do
       (r1cs, cs, _) = compileTExpToR1CS simpl texpPkg
   let [out_var] = r1cs_out_vars r1cs
   -- parse the inputs, either from cli or from file
-  (pubInputs, privInputs) <- do
+  (pubInputs, privInputs, outputs) <- do
     let assignmentsFP = mkAssignmentsFilePath assignments name
     eInput <- fromJSONLines <$> readFileLines assignmentsFP
     let inputAssignments = either (failWith . ErrMsg) id eInput
     pure $ splitInputAssignments inputAssignments
-  let out_interp = comp_interp comp pubInputs (Map.mapKeys fst privInputs)
-      witness@(Witness {witness_assgn = Assgn m}) = wit_of_cs pubInputs (Map.mapKeys snd privInputs) cs
+  let mOut = case outputs of
+        [(v, k)] ->
+          if v /= out_var
+            then failWith $ ErrMsg $ "derived output variable " <> show out_var <> " does not match provided output variable " <> show v
+            else Just (v, k)
+        [] -> Nothing
+        _ -> failWith $ ErrMsg $ "expected one output variable, got " <> show outputs
+      initAssignments = Map.mapKeys snd privInputs <> maybe mempty (uncurry Map.singleton) mOut
+  let witness@(Witness {witness_assgn = Assgn m}) = wit_of_cs pubInputs initAssignments cs
       out = case Map.lookup out_var m of
         Nothing ->
           failWith $
@@ -83,6 +90,7 @@ runAll RunAllOpts {..} name comp = do
                   ++ show witness
               )
         Just out_val -> out_val
+      out_interp = comp_interp comp pubInputs (Map.mapKeys fst privInputs)
   unless (out_interp == out) $
     failWith $
       ErrMsg $

src/Snarkl/Common.hs

@@ -161,6 +161,7 @@ instance A.FromJSON ConstraintHeader where
 data InputVar
   = PublicInputVar Var
   | PrivateInputVar String Var
+  | OutputVar Var
 
 instance A.ToJSON InputVar where
   toJSON (PublicInputVar v) =
@@ -174,32 +175,42 @@ instance A.ToJSON InputVar where
         "name" A..= name,
         "var" A..= v
       ]
+  toJSON (OutputVar v) =
+    A.object
+      [ "tag" A..= ("output" :: String),
+        "var" A..= v
+      ]
 
-splitInputVars :: [InputVar] -> ([Var], Map.Map String Var)
+splitInputVars :: [InputVar] -> ([Var], Map.Map String Var, [Var])
 splitInputVars =
   foldr
-    ( \iv (pubs, privs) ->
+    ( \iv (pubs, privs, outputs) ->
         case iv of
-          PublicInputVar v -> (v : pubs, privs)
-          PrivateInputVar name v -> (pubs, Map.insert name v privs)
+          PublicInputVar v -> (v : pubs, privs, outputs)
+          PrivateInputVar name v -> (pubs, Map.insert name v privs, outputs)
+          OutputVar v -> (pubs, privs, v : outputs)
     )
-    ([], Map.empty)
+    ([], Map.empty, [])
 
 instance A.FromJSON InputVar where
   parseJSON = A.withObject "InputVar" $ \v -> do
     tag <- v A..: "tag"
     case tag of
       ("public" :: String) -> PublicInputVar <$> v A..: "var"
       ("private" :: String) -> PrivateInputVar <$> v A..: "name" <*> v A..: "var"
+      ("output" :: String) -> OutputVar <$> v A..: "var"
       _ -> fail $ "unknown tag: " <> tag
 
 data InputAssignment k
   = PublicInputAssignment Var k
   | PrivateInputAssignment String Var k
+  | OutputAssignment Var k
+  deriving (Show)
 
 instance Functor InputAssignment where
   fmap f (PublicInputAssignment v k) = PublicInputAssignment v (f k)
   fmap f (PrivateInputAssignment name v k) = PrivateInputAssignment name v (f k)
+  fmap f (OutputAssignment v k) = OutputAssignment v (f k)
 
 instance (PrimeField k) => A.ToJSON (InputAssignment k) where
   toJSON (PublicInputAssignment v k) =
@@ -215,21 +226,29 @@ instance (PrimeField k) => A.ToJSON (InputAssignment k) where
         "var" A..= v,
         "value" A..= FieldElem k
       ]
+  toJSON (OutputAssignment v k) =
+    A.object
+      [ "tag" A..= ("output" :: String),
+        "var" A..= v,
+        "value" A..= FieldElem k
+      ]
 
 instance (PrimeField k) => A.FromJSON (InputAssignment k) where
   parseJSON = A.withObject "InputAssignment" $ \v -> do
     tag <- v A..: "tag"
     case tag of
       ("public" :: String) -> PublicInputAssignment <$> v A..: "var" <*> (unFieldElem <$> v A..: "value")
       ("private" :: String) -> PrivateInputAssignment <$> v A..: "name" <*> v A..: "var" <*> (unFieldElem <$> v A..: "value")
+      ("output" :: String) -> OutputAssignment <$> v A..: "var" <*> (unFieldElem <$> v A..: "value")
       _ -> fail $ "unknown tag: " <> tag
 
-splitInputAssignments :: [InputAssignment k] -> ([k], Map.Map (String, Var) k)
+splitInputAssignments :: [InputAssignment k] -> ([k], Map.Map (String, Var) k, [(Var, k)])
 splitInputAssignments =
   foldr
-    ( \ia (pubs, privs) ->
+    ( \ia (pubs, privs, outputs) ->
         case ia of
-          PublicInputAssignment _ k -> (k : pubs, privs)
-          PrivateInputAssignment name v k -> (pubs, Map.insert (name, v) k privs)
+          PublicInputAssignment _ k -> (k : pubs, privs, outputs)
+          PrivateInputAssignment name v k -> (pubs, Map.insert (name, v) k privs, outputs)
+          OutputAssignment var val -> (pubs, privs, (var, val) : outputs)
     )
-    ([], Map.empty)
+    ([], Map.empty, [])

tests/Test/ArkworksBridge.hs

@@ -4,13 +4,17 @@ import qualified Data.ByteString.Lazy as LBS
 import Data.Field.Galois (GaloisField, PrimeField)
 import Data.JSONLines (ToJSONLines (toJSONLines))
 import qualified Data.Map as Map
+import Data.Maybe (fromJust)
 import Data.Typeable (Typeable)
+import Debug.Trace (trace)
 import Snarkl.AST (Comp)
 import Snarkl.Backend.R1CS
 import Snarkl.Backend.R1CS.R1CS (witnessInputs)
 import Snarkl.CLI.Common (mkInputsFilePath, mkR1CSFilePath, mkWitnessFilePath)
-import Snarkl.Common (Assgn (Assgn), FieldElem (FieldElem))
+import Snarkl.Common (Assgn (Assgn), FieldElem (FieldElem), InputAssignment (OutputAssignment, PublicInputAssignment), Var (Var))
 import Snarkl.Compile (SimplParam, compileCompToR1CS)
+import Snarkl.Constraint (ConstraintSystem (cs_out_vars, cs_public_in_vars), SimplifiedConstraintSystem (unSimplifiedConstraintSystem))
+import Snarkl.Errors (ErrMsg (ErrMsg), failWith)
 import Snarkl.Toplevel (wit_of_cs)
 import qualified System.Exit as GHC
 import System.Process (createProcess, shell, waitForProcess)
@@ -37,6 +41,7 @@ runCMD (CreateTrustedSetup rootDir name simpl c) = do
   waitForProcess hdl
 runCMD (CreateProof rootDir name simpl c inputs) = do
   let (r1cs, simplifiedCS, _) = compileCompToR1CS simpl c
+  let public_in_vars = cs_public_in_vars (unSimplifiedConstraintSystem simplifiedCS)
       witness = wit_of_cs inputs Map.empty simplifiedCS
       r1csFilePath = mkR1CSFilePath rootDir name
       witsFilePath = mkWitnessFilePath rootDir name
@@ -54,13 +59,28 @@ runCMD (CreateProof rootDir name simpl c inputs) = do
   waitForProcess hdl
 runCMD (RunR1CS rootDir name simpl c inputs) = do
   let (r1cs, simplifiedCS, _) = compileCompToR1CS simpl c
+      [out] = cs_out_vars (unSimplifiedConstraintSystem simplifiedCS)
       wit@(Witness {witness_assgn = Assgn m}) = wit_of_cs inputs Map.empty simplifiedCS
-      r1csFilePath = mkR1CSFilePath rootDir name
+      outVal = case Map.lookup out m of
+        Nothing ->
+          failWith $
+            ErrMsg
+              ( "output variable "
+                  ++ show out
+                  ++ "not mapped, in\n  "
+                  ++ show wit
+              )
+        Just v -> v
+  let r1csFilePath = mkR1CSFilePath rootDir name
       witsFilePath = mkWitnessFilePath rootDir name
       inputsFilePath = mkInputsFilePath rootDir name
   LBS.writeFile r1csFilePath $ toJSONLines r1cs
   LBS.writeFile witsFilePath $ toJSONLines wit
-  LBS.writeFile inputsFilePath $ toJSONLines $ witnessInputs wit
+  let is =
+        let ls = trace (show outVal) $ Map.toList $ Map.delete out m
+         in -- TODO (fix this hack)
+            zipWith PublicInputAssignment (Var <$> [1 ..]) inputs <> [OutputAssignment out outVal]
+  LBS.writeFile inputsFilePath $ toJSONLines is
   let cmd =
         mkCommand
           "run-r1cs"

tests/Test/Snarkl/UnitSpec.hs

@@ -294,7 +294,7 @@ spec = do
       it "36-1" $ test_comp @F_BN128 [Simplify] prog36 [0] `shouldReturn` Right 10
       it "36-2" $ test_comp @F_BN128 [Simplify] prog36 [1] `shouldReturn` Right 7
 
-  describe "Keccak Tests" $ do
-    describe "keccak" $ do
-      it "keccak-2" $ test_comp @F_BN128 [Simplify] (keccak1 2) (fromIntegral <$> input_vals) `shouldReturn` Right 1
-      it "keccak-2" $ test_comp @F_BN128 [Simplify] (keccak1 5) (fromIntegral <$> input_vals) `shouldReturn` Right 1
+    describe "Keccak Tests" $
+      describe "keccak" $ do
+        it "keccak-2" $ test_comp @F_BN128 [Simplify] (keccak1 2) (fromIntegral <$> input_vals) `shouldReturn` Right 1
+        it "keccak-2" $ test_comp @F_BN128 [Simplify] (keccak1 5) (fromIntegral <$> input_vals) `shouldReturn` Right 1

tutorial/sudoku/solver/Main.hs

@@ -41,10 +41,10 @@ defaultMain ::
 defaultMain inputsFP assignmentsFP publicInputs = do
   putStrLn $ "Reading inputs from " <> inputsFP
   eInputs <- fromJSONLines <$> readFileLines inputsFP
-  let (_, privateInputs) = either error splitInputVars eInputs
+  let (_, privateInputs, [out]) = either error splitInputVars eInputs
   putStrLn "Solving puzzle..."
   m <- solvePuzzle publicInputs
-  let assignments = mkAssignments @F_BN128 publicInputs privateInputs m
+  let assignments = mkAssignments @F_BN128 publicInputs privateInputs m out
   writeFileWithDir assignmentsFP (toJSONLines $ NoHeader assignments)
   putStrLn $ "Writing assignments to " <> assignmentsFP
 
@@ -53,8 +53,9 @@ mkAssignments ::
   [Int] ->
   Map.Map String Var ->
   Map.Map (Int, Int) Int ->
+  Var ->
   [InputAssignment k]
-mkAssignments publicInputs privateInputs m =
+mkAssignments publicInputs privateInputs m out =
   let mkVarName (i, j) = "x_" <> show (i, j)
       privateAssignments =
         [ PrivateInputAssignment name var value
@@ -63,7 +64,7 @@ mkAssignments publicInputs privateInputs m =
             mkVarName pos == name
         ]
       publicAssignments = zipWith PublicInputAssignment (Var <$> [0 ..]) publicInputs
-   in fmap fromIntegral <$> publicAssignments <> privateAssignments
+   in fmap fromIntegral <$> publicAssignments <> privateAssignments <> [OutputAssignment out 1]
 
 mkPuzzleMap :: [Int] -> Map.Map (Int, Int) Int
 mkPuzzleMap xs = Map.fromList $ zip [(x, y) | x <- [0 .. 8], y <- [0 .. 8]] xs