1  
  2  {-# LANGUAGE BangPatterns, StrictData #-}
  3  module Slot where
  4  
  5  --------------------------------------------------------------------------------
  6  
  7  import Data.Bits
  8  import Data.Word
  9  import Data.Array
 10  
 11  import Data.ByteString (ByteString)
 12  import qualified Data.ByteString       as B
 13  import qualified Data.ByteString.Char8 as C
 14  
 15  import Control.Monad
 16  import System.IO
 17  
 18  import Poseidon2
 19  import Misc
 20  
 21  --------------------------------------------------------------------------------
 22  
 23  type Hash = Fr
 24  
 25  newtype Seed     = Seed     Int deriving (Eq,Show)
 26  newtype CellIdx  = CellIdx  Int deriving (Eq,Show)
 27  newtype BlockIdx = BlockIdx Int deriving (Eq,Show)
 28  newtype SlotIdx  = SlotIdx  Int deriving (Eq,Show)
 29  
 30  newtype CellData  = CellData  { fromCellData  :: ByteString }
 31  newtype BlockData = BlockData { fromBlockData :: ByteString }
 32  
 33  instance Show CellData  where show (CellData  bs) = "CellData<"  ++ show (B.length bs) ++ ">"
 34  instance Show BlockData where show (BlockData bs) = "BlockData<" ++ show (B.length bs) ++ ">"
 35  
 36  mkCellData :: SlotConfig -> ByteString -> CellData
 37  mkCellData cfg bs = if B.length bs == _cellSize cfg 
 38    then CellData bs
 39    else error $ "mkCellData: invalid cell size: " ++ show (B.length bs)
 40  
 41  mkBlockData :: SlotConfig -> ByteString -> BlockData
 42  mkBlockData cfg bs = if B.length bs == _blockSize cfg 
 43    then BlockData bs
 44    else error $ "mkBlockData: invalid block size: " ++ show (B.length bs)
 45  
 46  data DataSource
 47    = SlotFile FilePath
 48    | FakeData Seed
 49    deriving Show
 50  
 51  data SlotConfig = MkSlotCfg
 52    { _cellSize  :: Int           -- ^ cell size in bytes (eg. 2048)
 53    , _blockSize :: Int           -- ^ block size in bytes (eg. 65536)
 54    , _nCells    :: Int           -- ^ number of cells per slot (should be power of two)
 55    , _nSamples  :: Int           -- ^ how many cells we sample
 56    , _dataSrc   :: DataSource    -- ^ slot data source
 57    }
 58    deriving Show
 59  
 60  cellsPerBlock :: SlotConfig -> Int
 61  cellsPerBlock cfg = case divMod (_blockSize cfg) (_cellSize cfg) of
 62    (q,0) -> if q>1 then q else error "cells per block must be at least 2"
 63    _     -> error "block size is not divisible by the cell size"
 64  
 65  blocksPerSlot :: SlotConfig -> Int
 66  blocksPerSlot cfg = case divMod (_nCells cfg) (cellsPerBlock cfg) of
 67    (q,0) -> if q>1 then q else error "blocks per slot must be at least 2"
 68    _     -> error "slot size is not divisible by the block size"
 69  
 70  
 71  -- | Example slot configuration
 72  exSlotCfg :: SlotConfig
 73  exSlotCfg = MkSlotCfg
 74    { _cellSize  = 256
 75    , _blockSize = 4096
 76    , _nCells    = 1024
 77    , _nSamples  = 20
 78    , _dataSrc   = FakeData (Seed 12345)
 79    }
 80  
 81  fieldElemsPerCell :: SlotConfig -> Int
 82  fieldElemsPerCell cfg = (_cellSize cfg + 30) `div` 31
 83  
 84  --------------------------------------------------------------------------------
 85  -- * load data
 86  
 87  genFakeCell :: SlotConfig -> Seed -> CellIdx -> CellData
 88  genFakeCell cfg (Seed seed) (CellIdx idx) = (mkCellData cfg $ B.pack list) where
 89    list = go (fromIntegral $ _cellSize cfg) 1 
 90    seed1 = fromIntegral seed + 0xdeadcafe :: Word64
 91    seed2 = fromIntegral idx  + 0x98765432 :: Word64
 92    go :: Word64 -> Word64 -> [Word8]
 93    go 0   _     = []
 94    go cnt state = fromIntegral state'' : go (cnt-1) state'' where
 95      state'  = state*(state + seed1)*(state + seed2) + state*(state `xor` 0x5a5a5a5a) + seed1*state + (seed2 + 17)
 96      state'' = mod state' 1698428844001831
 97  
 98  genFakeBlock :: SlotConfig -> Seed -> BlockIdx -> BlockData
 99  genFakeBlock cfg seed (BlockIdx blockIdx) = (mkBlockData cfg $ B.concat$ map fromCellData cells) where
100    k = cellsPerBlock cfg
101    a =  k *  blockIdx
102    b =  k * (blockIdx + 1) - 1
103    cells = [ genFakeCell cfg seed (CellIdx j) | j<-[a..b] ]
104  
105  loadCellData :: SlotConfig -> CellIdx -> IO CellData
106  loadCellData cfg cellidx@(CellIdx idx) = case _dataSrc cfg of
107    FakeData seed  -> return $ genFakeCell cfg seed cellidx
108    SlotFile fname -> do
109      h <- openBinaryFile fname ReadMode
110      hSeek h AbsoluteSeek (fromIntegral (_cellSize cfg) * fromIntegral idx)
111      bs <- B.hGet h (_cellSize cfg)
112      hClose h
113      return (mkCellData cfg bs)
114  
115  loadBlockData :: SlotConfig -> BlockIdx -> IO BlockData
116  loadBlockData cfg blockidx@(BlockIdx idx) = case _dataSrc cfg of
117    FakeData seed  -> return $ genFakeBlock cfg seed blockidx
118    SlotFile fname -> do
119      h <- openBinaryFile fname ReadMode
120      hSeek h AbsoluteSeek (fromIntegral (_blockSize cfg) * fromIntegral idx)
121      bs <- B.hGet h (_blockSize cfg)
122      hClose h
123      return (mkBlockData cfg bs)
124  
125  --------------------------------------------------------------------------------
126  
127  {-
128  calcSlotTree :: SlotConfig -> IO MerkleTree
129  calcSlotTree cfg = calcMerkleTree <$> calcCellHashes cfg
130  
131  calcCellHashes :: SlotConfig -> IO [Hash]
132  calcCellHashes cfg = do
133    forM [0..(_nCells cfg - 1)] $ \idx -> do
134      cell <- loadCellData cfg idx 
135      return (hashCell cell)
136  -}
137  
138  --------------------------------------------------------------------------------
139  
140  -- | Split bytestring into smaller pieces, no padding
141  splitByteString :: Int -> ByteString -> [ByteString]
142  splitByteString k = go where
143    go bs 
144      | B.null bs  = []
145      | otherwise  = B.take k bs : go (B.drop k bs)
146  
147  splitBlockToCells :: SlotConfig -> BlockData -> [CellData]
148  splitBlockToCells cfg (BlockData blockdata) = 
149    map CellData (splitByteString (_cellSize cfg) blockdata)
150  
151  calcBlockTree :: SlotConfig -> BlockIdx -> IO MerkleTree
152  calcBlockTree cfg idx = do
153    block <- loadBlockData cfg idx
154    let cells = splitBlockToCells cfg block
155    let cellHashes = map (hashCell cfg) cells
156    let tree = calcMerkleTree cellHashes
157    return tree
158  
159  calcAllBlockTrees :: SlotConfig -> IO (Array Int MerkleTree)
160  calcAllBlockTrees cfg 
161    = listArray (0,n-1) <$> (forM [0..n-1] $ \idx -> calcBlockTree cfg (BlockIdx idx))
162    where
163      n = blocksPerSlot cfg
164  
165  --------------------------------------------------------------------------------
166  
167  data SlotTree = MkSlotTree
168    { _miniTrees :: Array Int MerkleTree     -- ^ block trees
169    , _bigTree   :: MerkleTree               -- ^ the tree over the block hashes
170    }
171  
172  slotTreeRoot :: SlotTree -> Hash
173  slotTreeRoot = merkleRootOf . _bigTree
174  
175  calcSlotTree :: SlotConfig -> IO SlotTree
176  calcSlotTree cfg = do
177    minitrees <- calcAllBlockTrees cfg
178    let bigtree = calcMerkleTree $ map merkleRootOf $ elems minitrees
179    return $ MkSlotTree minitrees bigtree
180  
181  extractCellProof :: SlotConfig -> SlotTree -> CellIdx -> [Hash]
182  extractCellProof cfg slotTree (CellIdx cellIdx) = final where
183    (blockIdx, withinBlockIdx) = cellIdx `divMod` (cellsPerBlock cfg)
184    blockTree = (_miniTrees slotTree) ! blockIdx
185    proof1 = extractMerkleProof blockTree           withinBlockIdx
186    proof2 = extractMerkleProof (_bigTree slotTree) blockIdx
187    final  = _merklePath proof1 ++ _merklePath proof2
188  
189  checkCellProof :: SlotConfig -> SlotTree -> CellIdx -> Hash -> [Hash] -> Bool
190  checkCellProof cfg slotTree (CellIdx cellIdx) cellHash path 
191    | logK + logM /= length path  = error "checkCellProof: incorrect Merkle path length"
192    | 2^logK /= k                 = error "checkCellProof: non-power-of-two number of cells per blocks"
193    | otherwise                   = reSlotHash == slotTreeRoot slotTree
194    where
195      k = cellsPerBlock cfg
196      m = blocksPerSlot cfg
197      logK = ceilingLog2 (fromIntegral k)
198      logM = ceilingLog2 (fromIntegral m)
199  
200      blockIdx       = shiftR cellIdx logK
201      inBlockCellIdx = cellIdx .&. (k-1)
202  
203      smallProof = MkMerkleProof
204        { _leafIndex   = inBlockCellIdx
205        , _leafData    = cellHash
206        , _merklePath  = take logK path
207        , _dataSize    = k
208        }
209      bigProof = MkMerkleProof
210        { _leafIndex   = blockIdx
211        , _leafData    = blockHash
212        , _merklePath  = drop logK path
213        , _dataSize    = m
214        }
215  
216      blockHash  = reconstructMerkleRoot smallProof  
217      reSlotHash = reconstructMerkleRoot bigProof
218  
219  --------------------------------------------------------------------------------
220  
221  -- | Hash a cell
222  hashCell :: SlotConfig -> CellData -> Hash
223  hashCell cfg (CellData rawdata) 
224    | B.length rawdata /= _cellSize cfg  = error "hashCell: invalid cell data size"
225    | otherwise                          = hashCell_ rawdata
226  
227  hashCell_ :: ByteString -> Hash
228  hashCell_ rawdata = sponge2 (cellDataToFieldElements $ CellData rawdata) 
229  
230  --------------------------------------------------------------------------------
231  
232  -- | A 31-byte long chunk
233  newtype Chunk 
234    = Chunk ByteString 
235    deriving Show
236  
237  -- | Split bytestring into samller pieces, applying the @10*@ padding strategy.
238  --
239  -- That is, always add a single @0x01@ byte, and then add the necessary
240  -- number (in the interval @[0..k-1]@) of @0x00@ bytes to be a multiple of the 
241  -- given chunk length
242  --
243  padAndSplitByteString :: Int -> ByteString -> [Chunk]
244  padAndSplitByteString k orig = go (B.snoc orig 0x01) where
245    go bs 
246      | m == 0      = []
247      | m < k       = [Chunk $ B.append bs (B.replicate (k-m) 0x00)]
248      | otherwise   = (Chunk $ B.take k bs) : go (B.drop k bs)
249      where
250        m = B.length bs
251  
252  -- | Chunk a ByteString into a sequence of field elements
253  cellDataToFieldElements :: CellData -> [Fr]
254  cellDataToFieldElements (CellData rawdata) = map chunkToField pieces where
255    chunkSize = 31
256    pieces = padAndSplitByteString chunkSize rawdata
257  
258  chunkToField :: Chunk -> Fr
259  chunkToField chunk@(Chunk bs)
260    | l == 31  = fromInteger (chunkToIntegerLE chunk)
261    | l <  31  = error "chunkToField: chunk is too small (expecting exactly 31 bytes)"
262    | l >  31  = error "chunkToField: chunk is too big (expecting exactly 31 bytes)"
263    where 
264      l = B.length bs
265  
266  -- | Interpret a ByteString as an integer (little-endian)
267  chunkToIntegerLE :: Chunk -> Integer
268  chunkToIntegerLE (Chunk chunk) = go (B.unpack chunk) where
269    go []     = 0
270    go (w:ws) = fromIntegral w + shiftL (go ws) 8
271  
272  --------------------------------------------------------------------------------