content: {-# LANGUAGE OverloadedStrings #-}
  module GenPrecomputed where
  
  import Prelude hiding (sum)
  
  import Data.List.Split (chunksOf)
  import Numeric (showHex)
  import Prettyprinter ( Doc, (<+>), encloseSep, flatAlt, fillSep, group, line, nest, pretty, vsep
                       , SimpleDocStream, LayoutOptions(..), PageWidth(..), defaultLayoutOptions, layoutPretty
                       , lbrace, rbrace, semi
                       )
  import Prettyprinter.Render.Text (renderIO)
  import System.IO (IOMode(WriteMode), withFile)
  
  import Simplicity.Digest
  import Simplicity.MerkleRoot
  import Simplicity.Programs.Generic
  import Simplicity.Tags
  import Simplicity.Ty
  import Simplicity.Ty.Word
  
  x <-> y = x <> line <> y
  
  class Enumerable a where
    allValues :: [a]
  
  instance Enumerable () where
    allValues = [()]
  
  instance (Enumerable a, Enumerable b) => Enumerable (Either a b) where
    allValues = (Left <$> allValues) ++ (Right <$> allValues)
  
  instance (Enumerable a, Enumerable b) => Enumerable (a, b) where
    allValues = pair <$> allValues <*> allValues
     where
      pair x y = (x, y)
  
  single = (:[])
  bracket = group . encloseSep (flatAlt "{ " "{") (flatAlt " }" "}") ", "
  
  prettyCHash :: Hash256 -> Doc a
  prettyCHash h = bracket (format <$> chunksOf 8 str_h)
   where
    format x = "0x" <> pretty x <> "u"
    str_h = padding ++ text
     where
      padding = replicate (64 - length text) '0'
      text = showHex (integerHash256 h) ""
  
  declIV :: String -> IV -> Doc a
  declIV name iv = nest 2 $ (pretty $ "static const sha256_midstate "++name++"IV =") <-> (bracket . single . prettyCHash $ ivHash iv) <> semi
  
  declareSignatureIV :: Doc a
  declareSignatureIV = vsep
                     [ "/* Initial values for Simplicity's standard tagged message digest. */"
                     , declIV "signature" signatureTag
                     ]
  
  declareTyIVs :: Doc a
  declareTyIVs = vsep $ "/* Initial values for all the 'typeName's. */":(declTy <$> ["unit", "sum", "prod"])
   where
    declTy name = declIV name (typeTag name)
  
  declareMRIVs :: Doc a
  declareMRIVs = vsep $ ["/* Initial values for all the tags for 'CMR's, 'AMR's and 'IMR's. */"]
              ++ (declCMR <$> names)
              ++ (declAMR <$> ["assertl", "assertr"] ++ names)
              ++ (declIMR <$> ["disconnect", "witness"])
              ++ [declIV "identity" identityHashTag, declIV "hidden" hiddenTag, declIV "jet" jetTag]
   where
    declCMR name = declIV ("cmr_" ++ name) (commitmentTag name)
    declAMR name = declIV ("amr_" ++ name) (annotatedTag name)
    declIMR name = declIV ("imr_" ++ name) (identityTag name)
    names = ["comp", "case", "pair", "disconnect", "injl", "injr", "take", "drop", "iden", "unit", "witness"]
  
  declareWordCMRs :: Doc a
  declareWordCMRs = vsep
                  [ "/* This array contains the cmr of all canonical expressions of type X |- TWO through X |- TWO^8 that output distinct values."
                  , " * word_cmr[0..1] contain the cmrs for scribe(0) .. scribe(1) : X | - TWO"
                  , " * word_cmr[2..5] contain the cmrs for scribe(0) .. scribe(3) : X | - TWO^2"
                  , " * word_cmr[6..21] contain the cmrs for scribe(0) .. scribe(15) : X | - TWO^4"
                  , " * word_cmr[22..277] contain the cmrs for scribe(0) .. scribe(255) : X | - TWO^8"
                  , " */"
                  ] <> line <> nest 2 (vsep
                  [ "static const sha256_midstate word_cmr[] ="
                  , bracket (concat [word1s, word2s, word4s, word8s])  <> semi
                  ])
   where
    word1s = [bracket . single . prettyCHash . commitmentRoot $ (scribe x :: CommitmentRoot () Word1) | x <- allValues]
    word2s = [bracket . single . prettyCHash . commitmentRoot $ (scribe x :: CommitmentRoot () Word2) | x <- allValues]
    word4s = [bracket . single . prettyCHash . commitmentRoot $ (scribe x :: CommitmentRoot () Word4) | x <- allValues]
    word8s = [bracket . single . prettyCHash . commitmentRoot $ (scribe x :: CommitmentRoot () Word8) | x <- allValues]
  
  allWordRoots :: [Hash256]
  allWordRoots = typeRoot one : iterate diag bit
   where
    bit = typeRoot (sum one one)
    diag x = ivHash $ compress (typeTag "prod") (x,x)
  
  declareWordTypeRoots :: Doc a
  declareWordTypeRoots = vsep
                  [ "/* word_type_root[0] contains the type root of the ONE type."
                  , " * word_type_root[1] contains the type root of the TWO type."
                  , " * word_type_root[2] contains the type root of the TWO^2 type."
                  , " * word_type_root[3] contains the type root of the TWO^4 type."
                  , " * ..."
                  , " * word_type_root[32] contains the type root of the TWO^(2^31) type."
                  , " */"
                  ] <> line <> nest 2 (vsep
                  [ "static const sha256_midstate word_type_root[] ="
                  , bracket [bracket . single . prettyCHash $ x | x <- take 33 allWordRoots] <> semi
                  ])
  
  header = vsep $
    [ "/* This file has been automatically generated by GenPecomputed.hs */"
    , "#ifndef SIMPLICITY_PRECOMPUTED_H"
    , "#define SIMPLICITY_PRECOMPUTED_H"
    , ""
    , "#include \"sha256.h\""
    ]
  
  footer = vsep $
    [ "#endif"
    ]
  
  precomputed_h :: SimpleDocStream a
  precomputed_h = layoutPretty layoutOptions $ vsep (map (<> line)
    [ header
    , declareSignatureIV
    , declareTyIVs
    , declareMRIVs
    , declareWordCMRs
    , declareWordTypeRoots
    , footer
    ])
   where
    layoutOptions = LayoutOptions { layoutPageWidth = AvailablePerLine 132 1 }
  
  renderFile name doc = withFile name WriteMode (\h -> renderIO h doc)
  
  main = do
    renderFile "precomputed.h" precomputed_h
  