1  
  2  {-# LANGUAGE BangPatterns, StrictData #-}
  3  module Sampling where
  4  
  5  --------------------------------------------------------------------------------
  6  
  7  import Control.Monad
  8  import System.IO
  9  
 10  import qualified Data.ByteString as B
 11  
 12  import Slot    as Slot
 13  import DataSet as DataSet
 14  import Poseidon2
 15  
 16  import qualified ZK.Algebra.Curves.BN128.Fr.Mont as Fr
 17  
 18  --------------------------------------------------------------------------------
 19  
 20  samplingTest :: DataSetCfg -> SlotIdx -> Entropy -> FilePath -> IO ()
 21  samplingTest dsetCfg slotIdx entropy fpath = do
 22    input <- calculateCircuitInput dsetCfg slotIdx entropy
 23    exportCircuitInput fpath input
 24  
 25  --------------------------------------------------------------------------------
 26  
 27  type Entropy = Fr
 28  
 29  -- | Given an entropy source, the slot root, and a counter, we compute a
 30  -- cell index to sample
 31  sampleCellIndex :: SlotConfig -> Entropy -> Hash -> Int -> CellIdx
 32  sampleCellIndex cfg entropy slotRoot counter = CellIdx (fromInteger idx) where
 33    u   = sponge2 [entropy , slotRoot , fromIntegral counter] :: Fr
 34    idx = (Fr.from u) `mod` n          :: Integer
 35    n   = (fromIntegral $ Slot._nCells cfg) :: Integer
 36   
 37  --------------------------------------------------------------------------------
 38  
 39  padWithZeros :: Int -> [Fr] -> [Fr]
 40  padWithZeros n xs 
 41    | m <= n     = xs ++ replicate (n-m) Fr.zero
 42    | otherwise  = error "padWithZeros: input too long"
 43    where
 44      m = length xs
 45  
 46  --------------------------------------------------------------------------------
 47  
 48  data CircuitInput = MkInput 
 49    { _entropy      :: Entropy       -- ^ public input
 50    , _dataSetRoot  :: Hash          -- ^ public input
 51    , _slotIndex    :: Int           -- ^ public input
 52    , _slotRoot     :: Hash          -- ^ private input
 53    , _slotProof    :: [Fr]          -- ^ private input
 54    , _slotsPerDSet :: Int           -- ^ private input
 55    , _cellsPerSlot :: Int           -- ^ private input
 56    , _cellData     :: [[Fr]]        -- ^ private input
 57    , _merklePaths  :: [[Fr]]        -- ^ private input
 58    }
 59    deriving Show
 60  
 61  -- | Calculate the the inputs for the storage proof circuit
 62  calculateCircuitInput :: DataSetCfg -> SlotIdx -> Entropy -> IO CircuitInput
 63  calculateCircuitInput dataSetCfg slotIdx@(SlotIdx sidx) entropy = do
 64    let nslots = _nSlots dataSetCfg
 65  
 66    let slotCfgs = [ dataSetSlotCfg dataSetCfg (SlotIdx i) | i <- [0..nslots-1] ]
 67    slotTrees <- mapM calcSlotTree slotCfgs 
 68    let !slotRoots = map slotTreeRoot slotTrees
 69    let !dsetTree  = calcMerkleTree slotRoots
 70    let !dsetRoot  = merkleRootOf dsetTree 
 71  
 72    let ourSlotCfg  = slotCfgs  !! sidx
 73    let ourSlotRoot = slotRoots !! sidx
 74    let ourSlotTree = slotTrees !! sidx
 75    let !idxs = [ sampleCellIndex ourSlotCfg entropy ourSlotRoot j | j <- [1..(Slot._nSamples ourSlotCfg)] ]
 76  
 77    cellData <- forM idxs $ \idx -> (cellDataToFieldElements <$> loadCellData ourSlotCfg idx)
 78    let !merklePaths = [ extractCellProof ourSlotCfg ourSlotTree idx | idx <- idxs ]
 79    return $ MkInput
 80      { _entropy      = entropy
 81      , _dataSetRoot  = dsetRoot
 82      , _slotIndex    = sidx
 83      , _slotRoot     = ourSlotRoot
 84      , _slotProof    = padWithZeros (_maxLog2NSlots dataSetCfg) $ extractMerkleProof_ dsetTree sidx
 85      , _slotsPerDSet = nslots
 86      , _cellsPerSlot = Slot._nCells ourSlotCfg
 87      , _cellData     = cellData
 88      , _merklePaths  = map (padWithZeros (_maxDepth dataSetCfg)) merklePaths
 89      }
 90  
 91  -- | Export the inputs of the storage proof circuits in JSON format,
 92  -- which @circom@ can consume.
 93  --
 94  -- NOTE: large numbers (field elements) must be encoded as JSON strings,
 95  -- not numbers, as Javascript cannot handle large numbers!
 96  --
 97  exportCircuitInput :: FilePath -> CircuitInput -> IO ()
 98  exportCircuitInput fpath input = do
 99    h <- openFile fpath WriteMode
100    hPutStrLn h $ "{ \"entropy\":          " ++ show (show (_entropy      input))
101    hPutStrLn h $ ", \"dataSetRoot\":      " ++ show (show (_dataSetRoot  input))
102    hPutStrLn h $ ", \"slotIndex\":        " ++ show (show (_slotIndex    input))
103    hPutStrLn h $ ", \"slotRoot\":         " ++ show (show (_slotRoot     input))
104    hPutStrLn h $ ", \"nSlotsPerDataSet\": " ++ show (show (_slotsPerDSet input))
105    hPutStrLn h $ ", \"nCellsPerSlot\":    " ++ show (show (_cellsPerSlot input))
106    hPutStrLn h $ ", \"slotProof\":"
107    hPrintList h 4 (map show $ _slotProof input)
108    hPutStrLn h $ ", \"cellData\":" 
109    hPrintListOfLists h ((map.map) show $ _cellData input)
110    hPutStrLn h $ ", \"merklePaths\": " 
111    hPrintListOfLists h ((map.map) show $ _merklePaths input)
112    hPutStrLn h $ "}"
113    hClose h
114  
115  --------------------------------------------------------------------------------
116  
117  trueFalses :: [Bool]
118  trueFalses = True : repeat False
119  
120  indent :: Int -> String
121  indent k = replicate k ' '
122  
123  hPrintList' :: Show a => Handle -> (Bool -> String) -> [a] -> IO ()
124  hPrintList' h indentation xs = do
125    forM_ (zip trueFalses xs) $ \(b,x) -> do
126      hPutStrLn h (indentation b ++ (if b then "[ " else ", ") ++ show x)
127    hPutStrLn h (indentation False ++ "]")
128  
129  hPrintList :: Show a => Handle -> Int -> [a] -> IO ()
130  hPrintList h indentBy xs = hPrintList' h (\_ -> indent indentBy) $ xs
131  
132  hPrintListOfLists :: Show a => Handle -> [[a]] -> IO ()
133  hPrintListOfLists h xss = 
134    do
135      forM_ (zip trueFalses xss) $ \(b,xs) -> hPrintList' h (myIndentation b) xs
136      hPutStrLn h ("    ]")
137    where
138      myIndentation True  True  = "    [ "
139      myIndentation False True  = "    , "
140      myIndentation _     False = "      "
141  
142  --------------------------------------------------------------------------------