From 8fcef5c1c4ab145d293377ee21855415b41bf6af Mon Sep 17 00:00:00 2001
From: Balazs Komuves <bkomuves@gmail.com>
Date: Fri, 14 Mar 2025 17:24:05 +0100
Subject: [PATCH] export the witness (Haskell)

---
 README.md             |  2 +-
 haskell/src/Export.hs | 67 +++++++++++++++++++++++++++++++++++++++++++
 haskell/src/Main.hs   | 14 ++++++---
 3 files changed, 78 insertions(+), 5 deletions(-)
 create mode 100644 haskell/src/Export.hs

diff --git a/README.md b/README.md
index 7df50b8..eceeae4 100644
--- a/README.md
+++ b/README.md
@@ -22,7 +22,7 @@ Compiler (in Haskell):
 - [x] parsing the graph file
 - [x] parsing json input
 - [x] naive interpreter
-- [ ] exporting the witness
+- [x] exporting the witness
 - [ ] constantine backend
 - [ ] zikkurat backend
 - [ ] arkworks backend
diff --git a/haskell/src/Export.hs b/haskell/src/Export.hs
new file mode 100644
index 0000000..afb111b
--- /dev/null
+++ b/haskell/src/Export.hs
@@ -0,0 +1,67 @@
+
+-- | Exporting the witness
+
+{-# LANGUAGE NumericUnderscores, RecordWildCards #-}
+module Export where
+
+--------------------------------------------------------------------------------
+
+import Data.Bits
+import Data.Word
+import Data.Array
+
+import Control.Monad
+import Data.Binary.Put
+
+import qualified Data.ByteString.Lazy as L
+
+import BN254 ( F , fromF , fieldPrime )
+import Witness
+
+--------------------------------------------------------------------------------
+
+exportWitness :: FilePath -> Witness -> IO ()
+exportWitness fpath ws = do
+  let bs = runPut (putWitness ws)
+  L.writeFile fpath bs
+
+--------------------------------------------------------------------------------
+
+putHeader :: Int -> Put
+putHeader witnessLen = do
+  -- global header
+  putWord32be 0x_77_74_6e_73                   -- magic word @"wtns"@
+  putWord32le 2                                -- version
+  putWord32le 2                                -- number of sections
+
+  -- section 1
+  putWord32le 1                                -- section id
+  putWord64le 0x28                             -- section length
+  putWord32le 32                               -- 32 bytes per field element
+  putInteger32LE fieldPrime                    -- the field prime
+  putWord32le (fromIntegral witnessLen)        -- number of witness elements
+
+  -- section 2
+  putWord32le 2                                -- section id
+  putWord64le (32 * fromIntegral witnessLen)   -- section length
+
+--------------------------------------------------------------------------------
+
+arrayLength :: Array Int a -> Int
+arrayLength arr = let (a,b) = bounds arr in (b-a+1)
+
+putWitness :: Array Int F -> Put
+putWitness xs = do 
+  putHeader (arrayLength xs)
+  putRawWitness (elems xs)
+
+putRawWitness :: [F] -> Put
+putRawWitness list = mapM_ putF list
+
+putF :: F -> Put
+putF x = putInteger32LE (fromF x)
+
+putInteger32LE :: Integer -> Put
+putInteger32LE x = forM_ [0..31] $ \k -> putWord8 (fromIntegral (shiftR x (k*8) .&. 0xff))
+
+--------------------------------------------------------------------------------
diff --git a/haskell/src/Main.hs b/haskell/src/Main.hs
index f27f58f..552ebd8 100644
--- a/haskell/src/Main.hs
+++ b/haskell/src/Main.hs
@@ -10,6 +10,7 @@ import Witness
 import Parser
 import Graph
 import JSON
+import Export
 
 --------------------------------------------------------------------------------
 {-
@@ -28,12 +29,17 @@ testInputs = Map.fromList
 -}
 --------------------------------------------------------------------------------
 
+graphFile = "../../tmp/graph4.bin"
+inputFile = "../../tmp/input4.json"
+wtnsFile  = "../../tmp/my4.wtns"
+
 main :: IO ()
 main = do
-  Right graph <- parseGraphFile "../../tmp/graph4.bin"
+  Right graph <- parseGraphFile graphFile
   putStrLn ""
-  inputs <- loadInputJsonFile "../../tmp/input4.json"
+  inputs <- loadInputJsonFile inputFile
   print (inputSignals $ graphMeta graph)
   let wtns = witnessCalc inputs graph
-  putStrLn ""
-  print wtns
\ No newline at end of file
+  -- putStrLn ""
+  -- print wtns
+  exportWitness wtnsFile wtns