-- | Generate test cases for Nim

module TestGen.TestSponge where

--------------------------------------------------------------------------------

import Data.Array
import Data.List
import Data.Word

import System.IO

import Goldilocks
import Sponge

--------------------------------------------------------------------------------

nimShowF :: F -> String
nimShowF x = "toF( " ++ show x ++ "'u64 )"

nimShowDigest :: Digest -> String
nimShowDigest (MkDigest a b c d) = "[ " ++ intercalate ", " (map nimShowF [a,b,c,d]) ++ " ]"

nimShowPair :: (Integer,Digest) -> String
nimShowPair (n,d) = "( " ++ show n ++ " , " ++ nimShowDigest d ++ " )"

showListWith :: (a -> String) -> [a] -> [String]
showListWith f xys = zipWith (++) prefix (map f xys) where
  prefix = "  [ " : repeat "  , "

----------------------------------------

digests :: String -> (Integer -> Digest) -> [Integer] -> String
digests varname f xs = unlines (header : stuff ++ footer) where
  header = "const " ++ varname ++ "* : array[" ++ show length xs ++ ", tuple[n:int,digest:F4]] = "
  footer = ["  ]",""]
  stuff  = showListWith nimShowPair [ (x, f x) | x<-xs ]

--------------------------------------------------------------------------------

feltDigest :: Rate -> Integer -> Digest
feltDigest rate max = hashFieldElems' rate $ (map fromInteger [1..max] :: [F])

byteDigest :: Rate -> Integer -> Digest
byteDigest rate max = hashBytes' rate $ (map fromInteger [1..max] :: [Word8])

--------------------------------------------------------------------------------

printTests :: IO ()
printTests = hPrintTests stdout

hPrintTests :: Handle -> IO ()
hPrintTests h = hPutStrLn h $ unlines $ 
  [ digests ("testcases_field_rate" ++ show r) (feltDigest (Rate r)) [0..80] | r<-[1..8] ] ++
  [ digests ("testcases_bytes_rate" ++ show r) (byteDigest (Rate r)) [0..80] | r<-[4,8]  ]

writeTests :: IO ()
writeTests = withFile "spongeTestCases.nim" WriteMode $ \h -> do
  hPutStrLn   h "# generated by TestGen/TestSponge.hs\n"
  hPutStrLn   h "import poseidon2/types\n"
  hPrintTests h

--------------------------------------------------------------------------------