GenRustJets.hs

{-# LANGUAGE OverloadedStrings #-}
module GenRustJets where

import Data.Char (toLower)
import Data.Foldable (toList)
import Data.Function (on)
import Data.Functor.Fixedpoint (Fix(..))
import Data.List (sortBy)
import Data.List.Split (chunksOf)
import Data.Maybe (fromMaybe)
import qualified Data.Map as Map
import Numeric (showHex)
import Prettyprinter ( Doc, (<+>), braces, comma, dquotes, fillSep, line, nest, parens, pretty, punctuate, semi, tupled, vsep
                     , SimpleDocStream, LayoutOptions(..), PageWidth(..), defaultLayoutOptions, layoutPretty
                     )
import Prettyprinter.Render.Text (renderIO)
import System.IO (IOMode(WriteMode), withFile)

import NameWrangler
import qualified Simplicity.Bitcoin.Jets as Bitcoin
import qualified Simplicity.Bitcoin.Term as Bitcoin
import qualified Simplicity.CoreJets as Core
import Simplicity.CoreJets (CoreJet)
import Simplicity.Digest
import qualified Simplicity.Elements.Jets as Elements
import qualified Simplicity.Elements.Term as Elements
import Simplicity.MerkleRoot
import Simplicity.Serialization
import Simplicity.Tree
import Simplicity.Ty
import Simplicity.Weight

x <-> y = x <> line <> y

nestBraces x = braces (nest 4 (line <> x) <> line)

data JetModule = CoreModule | BitcoinModule | ElementsModule
  deriving Eq

data JetData x y = JetData { jetName :: String
                           , jetCMR :: CommitmentRoot x y
                           , jetModule :: JetModule
                           , jetCost :: Weight
                           }

sortJetName = sortBy (compare `on` name)
 where
  name (SomeArrow j) = jetName j

cJetName = lowerSnakeCase . jetName

coreJetData :: (TyC x, TyC y) => CoreJet x y -> JetData x y
coreJetData jet = JetData { jetName = mkName jet
                          , jetCMR = cmr
                          , jetModule = CoreModule
                          , jetCost = cost
                          }
  where
    cmr = Bitcoin.asJet (Bitcoin.CoreJet jet)
    cost = Bitcoin.jetCost (Bitcoin.CoreJet jet)

elementsJetData :: (TyC x, TyC y) => Elements.JetType x y -> JetData x y
elementsJetData jet = JetData { jetName = mkName jet
                              , jetCMR = Elements.asJet jet
                              , jetModule = jetModule
                              , jetCost = Elements.jetCost jet
                              }
 where
  jetModule | Elements.CoreJet _ <- jet = CoreModule
            | otherwise = ElementsModule

bitcoinJetData :: (TyC x, TyC y) => Bitcoin.JetType x y -> JetData x y
bitcoinJetData jet = JetData { jetName = mkName jet
                             , jetCMR = Bitcoin.asJet jet
                             , jetModule = jetModule
                             , jetCost = Bitcoin.jetCost jet
                             }
 where
  jetModule | Bitcoin.CoreJet _ <- jet = CoreModule
            | otherwise = BitcoinModule

data Module = Module { moduleName :: Maybe String
                     , moduleCodes :: BinTree (SomeArrow JetData)
                     }
moduleJets :: Module -> [SomeArrow JetData]
moduleJets = sortJetName . toList . moduleCodes

rustModuleName = fromMaybe "Core" . moduleName
lowerRustModuleName = map toLower . rustModuleName

coreModule :: Module
coreModule = Module Nothing (someArrowMap coreJetData <$> (treeEvalBitStream Core.getJetBit))

-- Take Right is used to drop the (infinite) branch of constant word jets.
takeRight (Branch _ r) = r

elementsModule :: Module
elementsModule = Module (Just "Elements") (someArrowMap elementsJetData <$> takeRight (treeEvalBitStream Elements.getJetBit))

bitcoinModule :: Module
bitcoinModule = Module (Just "Bitcoin") (someArrowMap bitcoinJetData <$> takeRight (treeEvalBitStream Bitcoin.getJetBit))

data CompactTy = CTyOne
               | CTyWord Int
               | CTyMaybe CompactTy
               | CTySum CompactTy CompactTy
               | CTyProd CompactTy CompactTy
     deriving (Eq, Ord)

compactTy :: Ty -> CompactTy
compactTy = memoCataTy go
 where
  go One = CTyOne
  go (Sum CTyOne CTyOne) = CTyWord 1
  go (Sum CTyOne y) = CTyMaybe y
  go (Sum x y) = CTySum x y
  go (Prod (CTyWord wx) (CTyWord wy)) | wx == wy = CTyWord (wx + wy)
  go (Prod x y) = CTyProd x y

compactRustName :: CompactTy -> ShowS
compactRustName s = showString "b\"" . go s . showString "\""
 where
  go CTyOne = showString "1"
  go (CTyWord 1) = showString "2"
  go (CTyWord 32) = showString "i"
  go (CTyWord 64) = showString "l"
  go (CTyWord 256) = showString "h"
  go (CTyWord n) | even n = let rec = go (CTyWord (n `div` 2)) in showString "*" . rec . rec
  go (CTyMaybe x) = showString "+1" . go x
  go (CTySum x y) = showString "+" . go x . go y
  go (CTyProd x y) = showString "*" . go x . go y

showRustHash :: Hash256 -> Doc a
showRustHash h = fillSep $ ((<> comma) . format <$> chunksOf 2 str_h)
 where
  format x = pretty $ "0x" ++ x
  str_h = padding ++ text
   where
    padding = replicate (64 - length text) '0'
    text = showHex (integerHash256 h) ""

rustJetCmr :: Module -> Doc a
rustJetCmr mod = vsep $
  [ nest 4 (vsep ("fn cmr(&self) -> Cmr {" :
-- Temporary if statement until Bitcoin Jets have weight costs assigned to them
-- See Haskell/Simplicity/Bitcoin/Jets.hs:  jetCost (BitcoinJet jt) = error "Simplicity.Bitcoin.Jets.jetCost: :TODO: Implement jets for Bitcoin and benchmark them."
   if Just "Bitcoin" == moduleName mod
   then ["unimplemented!(\"Bitcoin jet CMRs weights have not yet been implemented.\")"]
   else
    [ nest 4 (vsep ("let bytes = match self {" :
        map (<>comma)
        [ nest 4 (vsep
          [ pretty modname <> "::" <> pretty (jetName jet) <+> "=> ["
          , showRustHash (commitmentRoot (jetCMR jet))
          ]) <-> "]"
        | (SomeArrow jet) <- moduleJets mod
        ]))
    , "};"
    , mempty
    , "Cmr(Midstate(bytes))"
    ]))
  , "}"
  ]
 where
  modname = rustModuleName mod

rustJetTy fname getTy mod = vsep $
  [ nest 4 (vsep (pretty ("fn "++fname++"(&self) -> TypeName {") :
    [ nest 4 (vsep ("let name: &'static [u8] = match self {" :
        map (<>comma)
        [ pretty modname <> "::" <> pretty (jetName jet) <+> "=>" <+>
          pretty (compactRustName (compactTy (getTy j)) "")
        | j@(SomeArrow jet) <- moduleJets mod
        ]))
    , "};"
    , mempty
    , "TypeName(name)"
    ]))
  , "}"
  ]
 where
  modname = rustModuleName mod

rustJetSourceTy :: Module -> Doc a
rustJetSourceTy = rustJetTy "source_ty" (\(SomeArrow jet) -> unreflect (fst (reifyArrow jet)))

rustJetTargetTy :: Module -> Doc a
rustJetTargetTy = rustJetTy "target_ty" (\(SomeArrow jet) -> unreflect (snd (reifyArrow jet)))

rustJetPtr :: Module -> Doc a
rustJetPtr mod = vsep $
  [ nest 4 (vsep ("fn c_jet_ptr(&self) -> &dyn Fn(&mut CFrameItem, CFrameItem, &Self::CJetEnvironment) -> bool {" :
    if modname == "Bitcoin"
    then ["unimplemented!(\"Bitcoin jets have not yet been implemented.\")"]
    else [ nest 4 (vsep ("match self {" :
        map (<>comma)
        [ pretty modname <> "::" <> pretty (jetName jet) <+> "=>" <+>
          pretty ("&simplicity_sys::c_jets::jets_wrapper::"++cJetName jet)
        | SomeArrow jet <- moduleJets mod
        ]))
    , "}"
    ]))
  , "}"
  ]
 where
  modname = rustModuleName mod

rustJetEncode :: Module -> Doc a
rustJetEncode mod =
  "fn encode<W: Write>(&self, w: &mut BitWriter<W>) -> std::io::Result<usize>" <+>
  nestBraces ("let (n, len) = match self" <+>
    nestBraces (vsep (foldMapWithPath item (moduleCodes mod))) <> semi <-> line <> "w.write_bits_be(n, len)")
 where
  item path (SomeArrow jet) = [pretty (rustModuleName mod ++ "::" ++ jetName jet) <+> "=>"
                          <+> tupled [pretty (code path 0 :: Int), pretty (length path)] <> comma]
  code [] n = n
  code (b : l) n = code l (2*n + if b then 1 else 0)

rustJetDecode :: Module -> Doc a
rustJetDecode mod =
  "fn decode<I: Iterator<Item = u8>>(bits: &mut BitIter<I>) -> Result<Self, decode::Error>" <+>
  nestBraces ("decode_bits!(bits," <+> braces (docTree (moduleCodes mod)) <> ")")
 where
  docTree Dead = mempty
  docTree (Leaf (SomeArrow jet)) = pretty (rustModuleName mod ++ "::" ++ jetName jet)
  docTree (Branch l r) = nest 4
                       ( line <> ("0" <+> "=>" <+> braces (docTree l)) <> comma
                     <-> ("1" <+> "=>" <+> braces (docTree r))
                       ) <> line

rustJetCost :: Module -> Doc a
rustJetCost mod = vsep $
  [ nest 4 (vsep ("fn cost(&self) -> Cost {" :
    if modname == "Bitcoin"
    then ["unimplemented!(\"Unspecified cost of Bitcoin jets\")"]
    else [ nest 4 (vsep ("match self {" :
        map (<>comma)
        [ pretty modname <> "::" <> pretty (jetName jet) <+> "=>" <+>
          "Cost::from_milliweight(" <> (pretty . milliWeight $ jetCost jet) <> ")"
        | SomeArrow jet <- moduleJets mod
        ]))
    , "}"
    ]))
  , "}"
  ]
 where
  modname = rustModuleName mod

rustJetImpl :: Module -> Doc a
rustJetImpl mod = vsep $
  [ nest 4 (vsep $ punctuate line
    ["impl Jet for" <+> pretty modname <+> "{"
    , env
    , rustJetCmr mod
    , rustJetSourceTy mod
    , rustJetTargetTy mod
    , rustJetEncode mod
    , rustJetDecode mod
    , rustJetPtr mod
    , rustJetCost mod
    ])
  , "}"
  ]
 where
  modname = rustModuleName mod
  env = vsep
    [ pretty $ "type Environment = "++env++";"
    , pretty $ "type CJetEnvironment = "++cEnv++";"
    , ""
    , pretty $ "fn c_jet_env("++envArg++": &Self::Environment) -> &Self::CJetEnvironment {"
    , pretty $ "    "++envBody
    , "}"
    ]
   where
    env | Nothing <- moduleName mod = "()"
        | Just "Elements" == moduleName mod = "ElementsEnv<std::sync::Arc<elements::Transaction>>"
        | Just name <- moduleName mod = name ++ "Env"
    cEnv | Just "Elements" == moduleName mod = "CElementsTxEnv"
         | otherwise = "()"
    envArg | Just "Bitcoin" == moduleName mod = "_env"
           | otherwise = "env"
    envBody | Nothing == moduleName mod = "env"
            | Just "Bitcoin" == moduleName mod = "unimplemented!(\"Unspecified CJetEnvironment for Bitcoin jets\")"
            | otherwise = "env.c_tx_env()"

rustJetEnum :: Module -> Doc a
rustJetEnum mod = vsep
 [ "/// The" <+> pretty (rustModuleName mod) <+> "jet family."
 , "#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]"
 , nest 4 $ vsep $
   ("pub enum" <+> pretty (rustModuleName mod) <+> "{") :
   [ pretty (jetName jet) <> comma | (SomeArrow jet) <- moduleJets mod ]
 , "}"
 , ""
 , nest 4 $ vsep $
   ("impl" <+> pretty (rustModuleName mod) <+> "{") :
   ("/// Array of all" <+> pretty (rustModuleName mod) <+> "jets.") :
   [ nest 4 $ vsep $
       ("pub const ALL: [Self;" <+> pretty (length $ moduleJets mod) <> "] = [") :
       [ "Self::" <> (pretty $ jetName jet) <> comma | (SomeArrow jet) <- moduleJets mod ]
     , "];"
   ]
 , "}"
 ]

rustJetDisplay :: Module -> Doc a
rustJetDisplay mod =
  "impl fmt::Display for" <+> pretty modname <+>
    nestBraces ("fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result" <+>
      nestBraces ("match self" <+>
        nestBraces (vsep (
          map (<>comma)
            [ pretty modname <> "::" <> pretty (jetName jet) <+> "=> f.write_str" <> (parens . dquotes . pretty $ cJetName jet)
            | SomeArrow jet <- moduleJets mod
            ]))
      )
    )
 where
  modname = rustModuleName mod

rustJetFromStr :: Module -> Doc a
rustJetFromStr mod =
  "impl str::FromStr for" <+> pretty modname <+>
    nestBraces (vsep
    [ "type Err = crate::Error;"
    , mempty
    , ("fn from_str(s: &str) -> Result<Self, Self::Err>" <+>
        nestBraces ("match s" <+>
          nestBraces (vsep (
            map (<> comma)
              ([ dquotes (pretty (cJetName jet)) <+> "=> Ok" <> parens (pretty modname <> "::" <> pretty (jetName jet))
               | SomeArrow jet <- moduleJets mod
               ] ++ [ "x => Err(crate::Error::InvalidJetName(x.to_owned()))" ]
              )))
      ))
    ])
 where
  modname = rustModuleName mod

rustHeader :: Doc a
rustHeader = "/* This file has been automatically generated. */"

rustImports :: Module -> Doc a
rustImports mod = vsep (map (<> semi)
  ([ "use crate::jet::type_name::TypeName"
  , "use crate::jet::Jet"
  , "use crate::merkle::cmr::Cmr"
  , "#[allow(unused_imports)]"
  , "use crate::decode_bits"
  , "use crate::{decode, BitIter, BitWriter}"
  , "use crate::analysis::Cost"
  , "use hashes::sha256::Midstate"
  , "use simplicity_sys::CFrameItem"
  , "use std::io::Write"
  , "use std::{fmt, str}"
  ] ++ envImports))
 where
  envImports | Nothing == moduleName mod = []
             | Just "Bitcoin" == moduleName mod = ["use crate::jet::bitcoin::BitcoinEnv"]
             | Just name <- moduleName mod =
             [ pretty $ "use crate::jet::"++map toLower name++"::"++name++"Env"
             , pretty $ "use simplicity_sys::C"++name++"TxEnv"
             ]

rustJetDoc :: Module -> SimpleDocStream a
rustJetDoc mod = layoutPretty layoutOptions $ vsep (map (<> line)
  [ rustHeader
  , rustImports mod
  , rustJetEnum mod
  , rustJetImpl mod
  , rustJetDisplay mod
  , rustJetFromStr mod
  ])

rustFFIImports :: Doc a
rustFFIImports = vsep (map (<> semi)
  [ "use crate::ffi::c_void"
  , "use crate::{CElementsTxEnv, CFrameItem}"
  ])

rustFFISigs :: Module -> Doc a
rustFFISigs mod = vsep
  [ nest 4 $ vsep $
    "extern \"C\" {" :
    (declaration <$> moduleJets mod)
  , "}"
  ]
 where
  declaration (SomeArrow jet) = (<> semi) . vsep $ pretty <$>
    [ linkName
    , signature
    ]
   where
    linkName = "#[link_name = \"c_"++cJetName jet++"\"]"
    signature = "pub fn "++cJetName jet++"(dst: *mut CFrameItem, src: *const CFrameItem, env: *const "++envType++") -> bool"
    envType | CoreModule <- jetModule jet = "c_void"
            | ElementsModule <- jetModule jet = "CElementsTxEnv"

rustFFIDoc :: Module -> SimpleDocStream a
rustFFIDoc mod = layoutPretty layoutOptions $ vsep (map (<> line)
  [ rustHeader
  , rustFFIImports
  , rustFFISigs mod
  ])

rustWrapperImports :: Doc a
rustWrapperImports = vsep (map (<> semi)
  [ "use crate::{CElementsTxEnv, CFrameItem}"
  , "use super::elements_ffi"
  ])

rustWrappers :: Module -> Doc a
rustWrappers mod = vsep ((<> line) . wrapper <$> moduleJets mod)
 where
  wrapper (SomeArrow jet) = vsep
   [ nest 4 $ vsep
     [ pretty $ "pub fn "++cJetName jet++templateParam++"(dst: &mut CFrameItem, src: CFrameItem, "++envParam++") -> bool {"
     , pretty $ "unsafe { "++lowerRustModuleName mod++"_ffi::"++cJetName jet++"(dst, &src, "++envArg++") }"
     ]
   , "}"
   ]
   where
    templateParam | CoreModule <- jetModule jet = "<T>"
                  | otherwise = ""
    envParam | CoreModule <- jetModule jet = "_env: &T"
             | ElementsModule <- jetModule jet = "env: &CElementsTxEnv"
    envArg | CoreModule <- jetModule jet = "std::ptr::null()"
           | ElementsModule <- jetModule jet = "env"

rustWrapperDoc :: Module -> SimpleDocStream a
rustWrapperDoc mod = layoutPretty layoutOptions $ vsep (map (<> line)
  [ rustHeader
  , rustWrapperImports
  , rustWrappers mod
  ])

cWrapperImports :: Doc a
cWrapperImports = vsep
  [ "#include \"simplicity/elements/elementsJets.h\""
  , "#include \"simplicity/simplicity_assert.h\""
  , "#include \"wrapper.h\""
  ]

cWrappers :: Module -> Doc a
cWrappers mod = vsep (map wrapper $ moduleJets mod)
 where
  wrapper (SomeArrow jet) = pretty $ "WRAP_("++cJetName jet++")"

cWrapperDoc :: Module -> SimpleDocStream a
cWrapperDoc mod = layoutPretty layoutOptions $ vsep (map (<> line)
  [ rustHeader -- also works for C
  , cWrapperImports
  , cWrappers mod
  ])

renderFile name doc = withFile name WriteMode (\h -> renderIO h doc)

main = do
  renderFile "core.rs" (rustJetDoc coreModule)
  renderFile "elements.rs" (rustJetDoc elementsModule)
  renderFile "bitcoin.rs" (rustJetDoc bitcoinModule)
  renderFile "jets_ffi.rs" (rustFFIDoc elementsModule)
  renderFile "jets_wrapper.rs" (rustWrapperDoc elementsModule)
  renderFile "jets_wrapper.c" (cWrapperDoc elementsModule)

layoutOptions = LayoutOptions { layoutPageWidth = AvailablePerLine 100 1 }