content: module EVM.RLP where
  
  import EVM.Types
  import Data.Bits (shiftR)
  import Data.ByteString (ByteString)
  import Data.ByteString qualified as BS
  import Witch (into, unsafeInto)
  
  data RLP = BS ByteString | List [RLP] deriving Eq
  
  instance Show RLP where
    show (BS str) = show (ByteStringS str)
    show (List list) = show list
  
  slice :: Int -> Int -> ByteString -> ByteString
  slice offset size bs = BS.take size $ BS.drop offset bs
  
  -- helper function returning (the length of the prefix, the length of the content, isList boolean, optimal boolean)
  itemInfo :: ByteString -> (Int, Int, Bool, Bool)
  itemInfo bs | bs == mempty = (0, 0, False, False)
              | otherwise = case BS.head bs of
    x | 0 <= x && x < 128   -> (0, 1, False, True) -- directly encoded byte
    x | 128 <= x && x < 184 -> (1, into x - 128, False, (BS.length bs /= 2) || (127 < (BS.head $ BS.drop 1 bs))) -- short string
    x | 184 <= x && x < 192 -> (1 + pre, len, False, (len > 55) && BS.head (BS.drop 1 bs) /= 0) -- long string
      where pre = into $ x - 183
            -- TODO: unsafeInto fails: cabal run test -- -p 'rlp' --quickcheck-replay=413899
            len = fromIntegral $ word $ slice 1 pre bs
    x | 192 <= x && x < 248 -> (1, into $ x - 192, True, True) -- short list
    x                       -> (1 + pre, len, True, (len > 55) && BS.head (BS.drop 1 bs) /= 0) -- long list
      where pre = into $ x - 247
            -- TODO: unsafeInto fails: cabal run test -- -p 'rlp' --quickcheck-replay=146332
            len = fromIntegral $ word $ slice 1 pre bs
  
  rlpdecode :: ByteString -> Maybe RLP
  rlpdecode bs =
    if optimal && pre + len == BS.length bs then
      if isList then do
        items <- mapM (\(s, e) -> rlpdecode $ slice s e content) (rlplengths content 0 len)
        Just (List items)
      else
        Just (BS content)
    else Nothing
    where (pre, len, isList, optimal) = itemInfo bs
          content = BS.drop pre bs
  
  rlplengths :: ByteString -> Int -> Int -> [(Int,Int)]
  rlplengths bs acc top =
    if acc < top then
      let (pre, len, _, _) = itemInfo bs
      in (acc, pre + len) : rlplengths (BS.drop (pre + len) bs) (acc + pre + len) top
    else []
  
  rlpencode :: RLP -> ByteString
  rlpencode (BS bs) = if BS.length bs == 1 && BS.head bs < 128 then bs
                      else encodeLen 128 bs
  rlpencode (List items) = encodeLen 192 (mconcat $ map rlpencode items)
  
  encodeLen :: Int -> ByteString -> ByteString
  encodeLen offset bs =
    if BS.length bs <= 55 then prefix (BS.length bs) <> bs
    else prefix lenLen <> lenBytes <> bs
    where
      lenBytes = asBE $ BS.length bs
      prefix n = BS.singleton $ unsafeInto $ offset + n
      lenLen = BS.length lenBytes + 55
  
  rlpList :: [RLP] -> ByteString
  rlpList n = rlpencode $ List n
  
  octets :: W256 -> ByteString
  octets x =
    BS.pack $ dropWhile (== 0) [fromIntegral (shiftR x (8 * i)) | i <- reverse [0..31]]
  
  octetsFull :: Int -> W256 -> ByteString
  octetsFull n x =
    BS.pack $ [fromIntegral (shiftR x (8 * i)) | i <- reverse [0..n]]
  
  octets160 :: Addr -> ByteString
  octets160 x =
    BS.pack $ dropWhile (== 0) [fromIntegral (shiftR x (8 * i)) | i <- reverse [0..19]]
  
  rlpWord256 :: W256 -> RLP
  rlpWord256 0 = BS mempty
  rlpWord256 n = BS $ octets n
  
  rlpWordFull :: W256 -> RLP
  rlpWordFull = BS . octetsFull 31
  
  rlpAddrFull :: Addr -> RLP
  rlpAddrFull = BS . octetsFull 19 . into
  
  rlpWord160 :: Addr -> RLP
  rlpWord160 0 = BS mempty
  rlpWord160 n = BS $ octets160 n
  