way easier and more consistent jsonl serialization

src/Data/JSONLines.hs

@@ -28,7 +28,7 @@ instance Monoid JSONLine where
 data WithHeader :: Type -> Type -> Type where
   WithHeader :: hdr -> [item] -> WithHeader hdr item
 
-data NoHeader item = NoHeader [item]
+newtype NoHeader item = NoHeader [item]
 
 jsonlBuilder ::
   (Foldable t) =>

src/Snarkl/Backend/R1CS/R1CS.hs

@@ -1,19 +1,18 @@
 module Snarkl.Backend.R1CS.R1CS
   ( R1C (..),
     R1CS (..),
-    r1csHeader,
     Witness (..),
-    witnessHeader,
     witnessInputs,
-    WitnessHeader (..),
     sat_r1cs,
     num_constraints,
   )
 where
 
 import Control.Parallel.Strategies
 import qualified Data.Aeson as A
+import Data.Bifunctor (Bifunctor (second))
 import Data.Field.Galois (GaloisField (char, deg), PrimeField)
+import Data.JSONLines (FromJSONLines (fromJSONLines), ToJSONLines (toJSONLines), WithHeader (..))
 import qualified Data.Map as Map
 import Snarkl.Backend.R1CS.Poly
 import Snarkl.Common
@@ -55,16 +54,29 @@ data R1CS a = R1CS
   }
   deriving (Show)
 
-r1csHeader :: (GaloisField a) => R1CS a -> ConstraintHeader
-r1csHeader (cs :: R1CS a) =
-  ConstraintHeader
-    { field_characteristic = toInteger $ char (undefined :: a),
-      extension_degree = toInteger $ deg (undefined :: a),
-      n_constraints = num_constraints cs,
-      n_variables = r1cs_num_vars cs,
-      input_variables = r1cs_in_vars cs,
-      output_variables = r1cs_out_vars cs
-    }
+instance (PrimeField k) => ToJSONLines (R1CS k) where
+  toJSONLines cs = toJSONLines $ WithHeader (r1csHeader cs) (r1cs_clauses cs)
+    where
+      r1csHeader (_ :: R1CS a) =
+        ConstraintHeader
+          { field_characteristic = toInteger $ char (undefined :: a),
+            extension_degree = toInteger $ deg (undefined :: a),
+            n_constraints = num_constraints cs,
+            n_variables = r1cs_num_vars cs,
+            input_variables = r1cs_in_vars cs,
+            output_variables = r1cs_out_vars cs
+          }
+
+instance (PrimeField k) => FromJSONLines (R1CS k) where
+  fromJSONLines ls = do
+    WithHeader ConstraintHeader {..} cs <- fromJSONLines ls
+    pure
+      R1CS
+        { r1cs_clauses = cs,
+          r1cs_num_vars = fromIntegral n_variables,
+          r1cs_in_vars = input_variables,
+          r1cs_out_vars = output_variables
+        }
 
 data Witness k = Witness
   { witness_assgn :: Assgn k,
@@ -99,13 +111,30 @@ instance A.FromJSON WitnessHeader where
       num_vars <- v A..: "n_variables"
       pure $ WitnessHeader in_vars out_vars num_vars
 
-witnessHeader :: Witness k -> WitnessHeader
-witnessHeader (Witness {..}) =
-  WitnessHeader
-    { in_vars = witness_in_vars,
-      out_vars = witness_out_vars,
-      num_vars = witness_num_vars
-    }
+instance (PrimeField k) => ToJSONLines (Witness k) where
+  toJSONLines wit@(Witness {witness_assgn = Assgn m}) =
+    toJSONLines $
+      WithHeader
+        (witnessHeader wit)
+        (Map.toList (FieldElem <$> m))
+    where
+      witnessHeader (Witness {..}) =
+        WitnessHeader
+          { in_vars = witness_in_vars,
+            out_vars = witness_out_vars,
+            num_vars = witness_num_vars
+          }
+
+instance (PrimeField k) => FromJSONLines (Witness k) where
+  fromJSONLines ls = do
+    WithHeader WitnessHeader {..} assgn <- fromJSONLines ls
+    pure
+      Witness
+        { witness_assgn = Assgn $ Map.fromList (second unFieldElem <$> assgn),
+          witness_in_vars = in_vars,
+          witness_out_vars = out_vars,
+          witness_num_vars = num_vars
+        }
 
 witnessInputs :: Witness k -> Assgn k
 witnessInputs (Witness {witness_in_vars, witness_assgn = Assgn m}) =

src/Snarkl/CLI/Compile.hs

@@ -18,13 +18,12 @@ import Data.JSONLines (ToJSONLines (toJSONLines), WithHeader (..))
 import qualified Data.Map as Map
 import Data.Maybe (catMaybes)
 import qualified Data.Set as Set
-import qualified Data.String.Conversions as CS
 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.Backend.R1CS
 import Snarkl.CLI.Common (mkConstraintsFilePath, mkR1CSFilePath, writeFileWithDir)
 import Snarkl.Compile
-import Snarkl.Constraint (ConstraintSystem (cs_constraints, cs_out_vars), SimplifiedConstraintSystem (unSimplifiedConstraintSystem), constraintSystemHeader)
+import Snarkl.Constraint (ConstraintSystem (cs_constraints, cs_out_vars), SimplifiedConstraintSystem (unSimplifiedConstraintSystem))
 import Snarkl.Errors (ErrMsg (ErrMsg), failWith)
 import Snarkl.Language
 import Snarkl.Toplevel (comp_interp, wit_of_cs)
@@ -86,10 +85,8 @@ compile CompileOpts {..} name comp = do
       TExpPkg nv in_vars e = compileCompToTexp comp
       (r1cs, cs) = compileTExpToR1CS simpl (TExpPkg nv in_vars e)
   let r1csFP = mkR1CSFilePath r1csOutput name
-  writeFileWithDir r1csFP . toJSONLines $
-    WithHeader (r1csHeader r1cs) (r1cs_clauses r1cs)
+  writeFileWithDir r1csFP $ toJSONLines r1cs
   putStrLn $ "Wrote R1CS to " <> r1csFP
   let csFP = mkConstraintsFilePath constraintsOutput name
-  writeFileWithDir csFP . toJSONLines $
-    WithHeader (constraintSystemHeader cs) (Set.toList $ cs_constraints $ unSimplifiedConstraintSystem cs)
+  writeFileWithDir csFP $ toJSONLines cs
   putStrLn $ "Wrote constraints to " <> csFP

src/Snarkl/CLI/GenWitness.hs

@@ -13,16 +13,16 @@ import Control.Applicative ((<|>))
 import Control.Monad (unless)
 import qualified Data.Aeson as A
 import Data.Field.Galois (PrimeField)
-import Data.JSONLines (FromJSONLines (fromJSONLines), NoHeader (..), ToJSONLines (toJSONLines), WithHeader (..))
+import Data.JSONLines (FromJSONLines (fromJSONLines), ToJSONLines (toJSONLines))
+import Data.List (sortOn)
 import qualified Data.Map as Map
-import qualified Data.Set as Set
 import qualified Data.String.Conversions as CS
 import Data.Typeable (Typeable)
 import Options.Applicative (Parser, eitherReader, help, long, option, showDefault, strOption, value)
 import Snarkl.Backend.R1CS
 import Snarkl.CLI.Common (mkConstraintsFilePath, mkR1CSFilePath, mkWitnessFilePath, readFileLines, writeFileWithDir)
-import Snarkl.Common (Assgn (Assgn), ConstraintHeader (..), FieldElem (..))
-import Snarkl.Constraint (ConstraintSystem (..), SimplifiedConstraintSystem (SimplifiedConstraintSystem, unSimplifiedConstraintSystem))
+import Snarkl.Common (Assgn (Assgn))
+import Snarkl.Constraint (ConstraintSystem (..), SimplifiedConstraintSystem (unSimplifiedConstraintSystem))
 import Snarkl.Errors (ErrMsg (ErrMsg), failWith)
 import Snarkl.Language (Comp)
 import Snarkl.Toplevel (comp_interp, wit_of_cs)
@@ -89,23 +89,15 @@ genWitness GenWitnessOpts {..} name comp = do
   constraints <- do
     let csFP = mkConstraintsFilePath constraintsInput name
     eConstraints <- fromJSONLines <$> readFileLines csFP
-    (WithHeader ConstraintHeader {..} cs) <- either (failWith . ErrMsg) pure eConstraints
-    pure $
-      SimplifiedConstraintSystem $
-        ConstraintSystem
-          { cs_constraints = Set.fromList cs,
-            cs_num_vars = n_variables,
-            cs_in_vars = input_variables,
-            cs_out_vars = output_variables
-          }
+    either (failWith . ErrMsg) pure eConstraints
   let [out_var] = cs_out_vars (unSimplifiedConstraintSystem constraints)
   -- parse the inputs, either from cli or from file
   is <- case inputs of
     Explicit is -> pure $ map fromInteger is
     FromFile fp -> do
       eInput <- fromJSONLines <$> readFileLines fp
-      NoHeader input <- either (failWith . ErrMsg) pure eInput
-      pure $ map unFieldElem input
+      let Assgn input = either (failWith . ErrMsg) id eInput
+      pure . map snd . sortOn fst $ Map.toList input
   let out_interp = comp_interp comp is
       witness@(Witness {witness_assgn = Assgn m}) = wit_of_cs is constraints
       out = case Map.lookup out_var m of
@@ -127,25 +119,13 @@ genWitness GenWitnessOpts {..} name comp = do
           ++ show out
   r1cs <- do
     let r1csFP = mkR1CSFilePath r1csInput name
-    (WithHeader ConstraintHeader {..} items) <- do
-      eConstraints <- fromJSONLines <$> readFileLines r1csFP
-      either (failWith . ErrMsg) pure eConstraints
-    pure $
-      R1CS
-        { r1cs_clauses = items,
-          r1cs_num_vars = n_variables,
-          r1cs_in_vars = input_variables,
-          r1cs_out_vars = output_variables
-        }
-  let witnessAssgn@(Assgn assgn) = witness_assgn witness
+    eConstraints <- fromJSONLines <$> readFileLines r1csFP
+    either (failWith . ErrMsg) pure eConstraints
   unless (sat_r1cs witness r1cs) $
     failWith $
       ErrMsg $
-        "witness\n  "
-          ++ CS.cs (A.encode witnessAssgn)
-          ++ "\nfailed to satisfy R1CS\n  "
+        "witness failed to satisfy R1CS\n  "
           ++ CS.cs (A.encode $ r1cs_clauses r1cs)
   let witnessFP = mkWitnessFilePath witnessOutput name
-  writeFileWithDir witnessFP . toJSONLines $
-    WithHeader (witnessHeader witness) (Map.toList (FieldElem <$> assgn))
+  writeFileWithDir witnessFP $ toJSONLines witness
   putStrLn $ "Wrote witness to " <> witnessFP

src/Snarkl/CLI/RunAll.hs

@@ -5,17 +5,17 @@ module Snarkl.CLI.RunAll where
 import Control.Monad (unless)
 import qualified Data.Aeson as A
 import Data.Field.Galois (PrimeField)
-import Data.JSONLines (FromJSONLines (..), NoHeader (..), ToJSONLines (..), WithHeader (..))
+import Data.JSONLines (FromJSONLines (..), NoHeader (..), ToJSONLines (..))
 import qualified Data.Map as Map
 import Data.Maybe (catMaybes)
 import qualified Data.String.Conversions as CS
 import Data.Typeable (Typeable)
 import Options.Applicative (Parser, help, long, showDefault, strOption, value)
-import Snarkl.Backend.R1CS (R1CS (r1cs_clauses, r1cs_out_vars), Witness (Witness, witness_assgn), r1csHeader, sat_r1cs)
+import Snarkl.Backend.R1CS (R1CS (r1cs_clauses, r1cs_out_vars), Witness (Witness, witness_assgn), sat_r1cs)
 import Snarkl.CLI.Common (mkR1CSFilePath, mkWitnessFilePath, readFileLines, writeFileWithDir)
 import Snarkl.CLI.Compile (OptimizeOpts (removeUnreachable, simplify), optimizeOptsParser)
 import Snarkl.CLI.GenWitness (InputOpts (Explicit, FromFile), inputOptsParser)
-import Snarkl.Common (Assgn (Assgn), FieldElem (FieldElem, unFieldElem))
+import Snarkl.Common (Assgn (Assgn), unFieldElem)
 import Snarkl.Compile (SimplParam (RemoveUnreachable, Simplify), TExpPkg (TExpPkg), compileCompToTexp, compileTExpToR1CS)
 import Snarkl.Errors (ErrMsg (ErrMsg), failWith)
 import Snarkl.Language (Comp)
@@ -89,19 +89,15 @@ runAll RunAllOpts {..} name comp = do
           ++ show out_interp
           ++ " differs from actual result "
           ++ show out
-  let witnessAssgn@(Assgn assgn) = witness_assgn witness
   unless (sat_r1cs witness r1cs) $
     failWith $
       ErrMsg $
         "witness\n  "
-          ++ CS.cs (A.encode witnessAssgn)
           ++ "\nfailed to satisfy R1CS\n  "
           ++ CS.cs (A.encode $ r1cs_clauses r1cs)
   let r1csFP = mkR1CSFilePath r1csOutput name
-  writeFileWithDir r1csFP . toJSONLines $
-    WithHeader (r1csHeader r1cs) (r1cs_clauses r1cs)
+  writeFileWithDir r1csFP $ toJSONLines r1cs
   putStrLn $ "Wrote R1CS to " <> r1csFP
   let witnessFP = mkWitnessFilePath witnessOutput name
-  writeFileWithDir witnessFP . toJSONLines $
-    WithHeader (r1csHeader r1cs) (Map.toList $ fmap FieldElem assgn)
+  writeFileWithDir witnessFP $ toJSONLines witness
   putStrLn $ "Wrote witness to " <> witnessFP

src/Snarkl/Common.hs

@@ -3,8 +3,10 @@
 module Snarkl.Common where
 
 import qualified Data.Aeson as A
+import Data.Bifunctor (Bifunctor (second))
 import Data.Field.Galois (PrimeField (fromP))
-import Data.Foldable (Foldable (toList))
+import Data.Foldable (toList)
+import Data.JSONLines (FromJSONLines (fromJSONLines), NoHeader (..), ToJSONLines (toJSONLines))
 import qualified Data.Map as Map
 import Text.PrettyPrint.Leijen.Text (Pretty (pretty))
 
@@ -29,6 +31,18 @@ incVar (Var i) = Var (i + 1)
 
 newtype Assgn a = Assgn (Map.Map Var a) deriving (Show, Eq, Functor)
 
+instance (PrimeField a) => A.ToJSON (Assgn a) where
+  toJSON (Assgn m) =
+    let kvs = map (\(var, coeff) -> (FieldElem coeff, var)) $ Map.toList m
+     in A.toJSON kvs
+
+instance (PrimeField a) => A.FromJSON (Assgn a) where
+  parseJSON =
+    A.withArray "Assgn" $ \arr -> do
+      kvs <- mapM A.parseJSON arr
+      let m = Map.fromList $ map (\(FieldElem k, v) -> (v, k)) (toList kvs)
+      return (Assgn m)
+
 -- We use this wrapper to get a stringified representation of big integers.
 -- This plays better with downstream dependencies, e.g. javascript.
 newtype FieldElem k = FieldElem {unFieldElem :: k} deriving (Show, Eq)
@@ -41,17 +55,13 @@ instance (PrimeField k) => A.FromJSON (FieldElem k) where
     k <- A.parseJSON v
     return $ FieldElem (fromInteger $ read k)
 
-instance (PrimeField a) => A.ToJSON (Assgn a) where
-  toJSON (Assgn m) =
-    let kvs = map (\(var, coeff) -> (FieldElem coeff, var)) $ Map.toList m
-     in A.toJSON kvs
+instance (PrimeField k) => ToJSONLines (Assgn k) where
+  toJSONLines (Assgn m) = toJSONLines $ NoHeader (Map.toList $ FieldElem <$> m)
 
-instance (PrimeField a) => A.FromJSON (Assgn a) where
-  parseJSON =
-    A.withArray "Assgn" $ \arr -> do
-      kvs <- mapM A.parseJSON arr
-      let m = Map.fromList $ map (\(FieldElem k, v) -> (v, k)) (toList kvs)
-      return (Assgn m)
+instance (PrimeField k) => FromJSONLines (Assgn k) where
+  fromJSONLines ls = do
+    NoHeader kvs <- fromJSONLines ls
+    pure . Assgn . Map.fromList $ map (second unFieldElem) kvs
 
 data UnOp = ZEq
   deriving (Eq, Show)

src/Snarkl/Constraint/Constraints.hs

@@ -9,7 +9,6 @@ module Snarkl.Constraint.Constraints
     ConstraintSet,
     ConstraintSystem (..),
     SimplifiedConstraintSystem (..),
-    constraintSystemHeader,
     r1cs_of_cs,
     renumber_constraints,
     constraint_vars,
@@ -20,6 +19,7 @@ import Control.Monad.State (State)
 import qualified Data.Aeson as A
 import Data.Bifunctor (Bifunctor (first, second))
 import Data.Field.Galois (GaloisField (char, deg), Prime, PrimeField)
+import Data.JSONLines (FromJSONLines (fromJSONLines), ToJSONLines (toJSONLines), WithHeader (..))
 import qualified Data.Map as Map
 import qualified Data.Set as Set
 import Snarkl.Backend.R1CS (Poly (Poly), R1C (R1C), R1CS (R1CS), const_poly, var_poly)
@@ -220,19 +220,31 @@ newtype SimplifiedConstraintSystem a = SimplifiedConstraintSystem
   }
   deriving (Show, Eq)
 
-constraintSystemHeader ::
-  (GaloisField a) =>
-  SimplifiedConstraintSystem a ->
-  ConstraintHeader
-constraintSystemHeader (SimplifiedConstraintSystem (ConstraintSystem {..}) :: SimplifiedConstraintSystem a) =
-  ConstraintHeader
-    { field_characteristic = toInteger $ char (undefined :: a),
-      extension_degree = toInteger $ deg (undefined :: a),
-      n_constraints = Set.size cs_constraints,
-      n_variables = cs_num_vars,
-      input_variables = cs_in_vars,
-      output_variables = cs_out_vars
-    }
+instance (PrimeField k) => ToJSONLines (SimplifiedConstraintSystem k) where
+  toJSONLines scs@(SimplifiedConstraintSystem (ConstraintSystem {..})) =
+    toJSONLines $ WithHeader (constraintSystemHeader scs) (Set.toList cs_constraints)
+    where
+      constraintSystemHeader (_ :: SimplifiedConstraintSystem a) =
+        ConstraintHeader
+          { field_characteristic = toInteger $ char (undefined :: a),
+            extension_degree = toInteger $ deg (undefined :: a),
+            n_constraints = Set.size cs_constraints,
+            n_variables = cs_num_vars,
+            input_variables = cs_in_vars,
+            output_variables = cs_out_vars
+          }
+
+instance (PrimeField k) => FromJSONLines (SimplifiedConstraintSystem k) where
+  fromJSONLines ls = do
+    WithHeader ConstraintHeader {..} cs <- fromJSONLines ls
+    pure $
+      SimplifiedConstraintSystem $
+        ConstraintSystem
+          { cs_constraints = Set.fromList cs,
+            cs_num_vars = fromIntegral n_variables,
+            cs_in_vars = input_variables,
+            cs_out_vars = output_variables
+          }
 
 r1cs_of_cs ::
   (GaloisField a) =>