{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}

module EVM.Patricia where

import EVM.RLP
import EVM.Types hiding (Literal)

import Control.Monad.Free
import Control.Monad.State
import Data.ByteString (ByteString)
import Data.Foldable (toList)
import Data.List (stripPrefix)
import Data.Sequence (Seq)

import qualified Data.ByteString as BS
import qualified Data.Map as Map
import qualified Data.Sequence as Seq

data KV k v a
  = Put k v a
  | Get k (v -> a)
  deriving (a -> KV k v b -> KV k v a
(a -> b) -> KV k v a -> KV k v b
(forall a b. (a -> b) -> KV k v a -> KV k v b)
-> (forall a b. a -> KV k v b -> KV k v a) -> Functor (KV k v)
forall a b. a -> KV k v b -> KV k v a
forall a b. (a -> b) -> KV k v a -> KV k v b
forall k v a b. a -> KV k v b -> KV k v a
forall k v a b. (a -> b) -> KV k v a -> KV k v b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> KV k v b -> KV k v a
$c<$ :: forall k v a b. a -> KV k v b -> KV k v a
fmap :: (a -> b) -> KV k v a -> KV k v b
$cfmap :: forall k v a b. (a -> b) -> KV k v a -> KV k v b
Functor)

newtype DB k v a = DB (Free (KV k v) a)
  deriving (a -> DB k v b -> DB k v a
(a -> b) -> DB k v a -> DB k v b
(forall a b. (a -> b) -> DB k v a -> DB k v b)
-> (forall a b. a -> DB k v b -> DB k v a) -> Functor (DB k v)
forall a b. a -> DB k v b -> DB k v a
forall a b. (a -> b) -> DB k v a -> DB k v b
forall k v a b. a -> DB k v b -> DB k v a
forall k v a b. (a -> b) -> DB k v a -> DB k v b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: a -> DB k v b -> DB k v a
$c<$ :: forall k v a b. a -> DB k v b -> DB k v a
fmap :: (a -> b) -> DB k v a -> DB k v b
$cfmap :: forall k v a b. (a -> b) -> DB k v a -> DB k v b
Functor, Functor (DB k v)
a -> DB k v a
Functor (DB k v)
-> (forall a. a -> DB k v a)
-> (forall a b. DB k v (a -> b) -> DB k v a -> DB k v b)
-> (forall a b c.
    (a -> b -> c) -> DB k v a -> DB k v b -> DB k v c)
-> (forall a b. DB k v a -> DB k v b -> DB k v b)
-> (forall a b. DB k v a -> DB k v b -> DB k v a)
-> Applicative (DB k v)
DB k v a -> DB k v b -> DB k v b
DB k v a -> DB k v b -> DB k v a
DB k v (a -> b) -> DB k v a -> DB k v b
(a -> b -> c) -> DB k v a -> DB k v b -> DB k v c
forall a. a -> DB k v a
forall k v. Functor (DB k v)
forall a b. DB k v a -> DB k v b -> DB k v a
forall a b. DB k v a -> DB k v b -> DB k v b
forall a b. DB k v (a -> b) -> DB k v a -> DB k v b
forall k v a. a -> DB k v a
forall a b c. (a -> b -> c) -> DB k v a -> DB k v b -> DB k v c
forall k v a b. DB k v a -> DB k v b -> DB k v a
forall k v a b. DB k v a -> DB k v b -> DB k v b
forall k v a b. DB k v (a -> b) -> DB k v a -> DB k v b
forall k v a b c. (a -> b -> c) -> DB k v a -> DB k v b -> DB k v c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: DB k v a -> DB k v b -> DB k v a
$c<* :: forall k v a b. DB k v a -> DB k v b -> DB k v a
*> :: DB k v a -> DB k v b -> DB k v b
$c*> :: forall k v a b. DB k v a -> DB k v b -> DB k v b
liftA2 :: (a -> b -> c) -> DB k v a -> DB k v b -> DB k v c
$cliftA2 :: forall k v a b c. (a -> b -> c) -> DB k v a -> DB k v b -> DB k v c
<*> :: DB k v (a -> b) -> DB k v a -> DB k v b
$c<*> :: forall k v a b. DB k v (a -> b) -> DB k v a -> DB k v b
pure :: a -> DB k v a
$cpure :: forall k v a. a -> DB k v a
$cp1Applicative :: forall k v. Functor (DB k v)
Applicative, Applicative (DB k v)
a -> DB k v a
Applicative (DB k v)
-> (forall a b. DB k v a -> (a -> DB k v b) -> DB k v b)
-> (forall a b. DB k v a -> DB k v b -> DB k v b)
-> (forall a. a -> DB k v a)
-> Monad (DB k v)
DB k v a -> (a -> DB k v b) -> DB k v b
DB k v a -> DB k v b -> DB k v b
forall a. a -> DB k v a
forall k v. Applicative (DB k v)
forall a b. DB k v a -> DB k v b -> DB k v b
forall a b. DB k v a -> (a -> DB k v b) -> DB k v b
forall k v a. a -> DB k v a
forall k v a b. DB k v a -> DB k v b -> DB k v b
forall k v a b. DB k v a -> (a -> DB k v b) -> DB k v b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: a -> DB k v a
$creturn :: forall k v a. a -> DB k v a
>> :: DB k v a -> DB k v b -> DB k v b
$c>> :: forall k v a b. DB k v a -> DB k v b -> DB k v b
>>= :: DB k v a -> (a -> DB k v b) -> DB k v b
$c>>= :: forall k v a b. DB k v a -> (a -> DB k v b) -> DB k v b
$cp1Monad :: forall k v. Applicative (DB k v)
Monad)

insertDB :: k -> v -> DB k v ()
insertDB :: k -> v -> DB k v ()
insertDB k
k v
v = Free (KV k v) () -> DB k v ()
forall k v a. Free (KV k v) a -> DB k v a
DB (Free (KV k v) () -> DB k v ()) -> Free (KV k v) () -> DB k v ()
forall a b. (a -> b) -> a -> b
$ KV k v () -> Free (KV k v) ()
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (KV k v () -> Free (KV k v) ()) -> KV k v () -> Free (KV k v) ()
forall a b. (a -> b) -> a -> b
$ k -> v -> () -> KV k v ()
forall k v a. k -> v -> a -> KV k v a
Put k
k v
v ()

lookupDB :: k -> DB k v v
lookupDB :: k -> DB k v v
lookupDB k
k = Free (KV k v) v -> DB k v v
forall k v a. Free (KV k v) a -> DB k v a
DB (Free (KV k v) v -> DB k v v) -> Free (KV k v) v -> DB k v v
forall a b. (a -> b) -> a -> b
$ KV k v v -> Free (KV k v) v
forall (f :: * -> *) (m :: * -> *) a.
(Functor f, MonadFree f m) =>
f a -> m a
liftF (KV k v v -> Free (KV k v) v) -> KV k v v -> Free (KV k v) v
forall a b. (a -> b) -> a -> b
$ k -> (v -> v) -> KV k v v
forall k v a. k -> (v -> a) -> KV k v a
Get k
k v -> v
forall a. a -> a
id

-- Collapses a series of puts and gets down to the monad of your choice
runDB :: Monad m
      => (k -> v -> m ()) -- ^ The 'put' function for our desired monad
      -> (k -> m v)       -- ^ The 'get' function for the same monad
      -> DB k v a         -- ^ The puts and gets to execute
      -> m a
runDB :: (k -> v -> m ()) -> (k -> m v) -> DB k v a -> m a
runDB k -> v -> m ()
putt k -> m v
gett (DB Free (KV k v) a
ops) = Free (KV k v) a -> m a
go Free (KV k v) a
ops
  where
    go :: Free (KV k v) a -> m a
go (Pure a
a) = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a
    go (Free (Put k
k v
v Free (KV k v) a
next)) = k -> v -> m ()
putt k
k v
v m () -> m a -> m a
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Free (KV k v) a -> m a
go Free (KV k v) a
next
    go (Free (Get k
k v -> Free (KV k v) a
handler)) = k -> m v
gett k
k m v -> (v -> m a) -> m a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Free (KV k v) a -> m a
go (Free (KV k v) a -> m a) -> (v -> Free (KV k v) a) -> v -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Free (KV k v) a
handler

type Path = [Nibble]

data Ref = Hash ByteString | Literal Node
  deriving (Ref -> Ref -> Bool
(Ref -> Ref -> Bool) -> (Ref -> Ref -> Bool) -> Eq Ref
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Ref -> Ref -> Bool
$c/= :: Ref -> Ref -> Bool
== :: Ref -> Ref -> Bool
$c== :: Ref -> Ref -> Bool
Eq)

instance Show Ref where
  show :: Ref -> String
show (Hash ByteString
d) = ByteStringS -> String
forall a. Show a => a -> String
show (ByteString -> ByteStringS
ByteStringS ByteString
d)
  show (Literal Node
n) = Node -> String
forall a. Show a => a -> String
show Node
n

data Node = Empty
          | Shortcut Path (Either Ref ByteString)
          | Full (Seq Ref) ByteString
  deriving (Int -> Node -> ShowS
[Node] -> ShowS
Node -> String
(Int -> Node -> ShowS)
-> (Node -> String) -> ([Node] -> ShowS) -> Show Node
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Node] -> ShowS
$cshowList :: [Node] -> ShowS
show :: Node -> String
$cshow :: Node -> String
showsPrec :: Int -> Node -> ShowS
$cshowsPrec :: Int -> Node -> ShowS
Show, Node -> Node -> Bool
(Node -> Node -> Bool) -> (Node -> Node -> Bool) -> Eq Node
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Node -> Node -> Bool
$c/= :: Node -> Node -> Bool
== :: Node -> Node -> Bool
$c== :: Node -> Node -> Bool
Eq)

-- the function HP from Appendix C of yellow paper
encodePath :: Path -> Bool -> ByteString
encodePath :: Path -> Bool -> ByteString
encodePath Path
p Bool
isTerminal | Int -> Bool
forall a. Integral a => a -> Bool
even (Path -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Path
p)
  = Path -> ByteString
packNibbles (Path -> ByteString) -> Path -> ByteString
forall a b. (a -> b) -> a -> b
$ Word8 -> Nibble
Nibble Word8
flag Nibble -> Path -> Path
forall a. a -> [a] -> [a]
: Word8 -> Nibble
Nibble Word8
0 Nibble -> Path -> Path
forall a. a -> [a] -> [a]
: Path
p
                        | Bool
otherwise
  = Path -> ByteString
packNibbles (Path -> ByteString) -> Path -> ByteString
forall a b. (a -> b) -> a -> b
$ Word8 -> Nibble
Nibble (Word8
flag Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
1) Nibble -> Path -> Path
forall a. a -> [a] -> [a]
: Path
p
  where flag :: Word8
flag  = if Bool
isTerminal then Word8
2 else Word8
0

rlpRef :: Ref -> RLP
rlpRef :: Ref -> RLP
rlpRef (Hash ByteString
d) = ByteString -> RLP
BS ByteString
d
rlpRef (Literal Node
n) = Node -> RLP
rlpNode Node
n

rlpNode :: Node -> RLP
rlpNode :: Node -> RLP
rlpNode Node
Empty = ByteString -> RLP
BS ByteString
forall a. Monoid a => a
mempty
rlpNode (Shortcut Path
path (Right ByteString
val)) = [RLP] -> RLP
List [ByteString -> RLP
BS (ByteString -> RLP) -> ByteString -> RLP
forall a b. (a -> b) -> a -> b
$ Path -> Bool -> ByteString
encodePath Path
path Bool
True, ByteString -> RLP
BS ByteString
val]
rlpNode (Shortcut Path
path (Left  Ref
ref)) = [RLP] -> RLP
List [ByteString -> RLP
BS (ByteString -> RLP) -> ByteString -> RLP
forall a b. (a -> b) -> a -> b
$ Path -> Bool -> ByteString
encodePath Path
path Bool
False, Ref -> RLP
rlpRef Ref
ref]
rlpNode (Full Seq Ref
refs ByteString
val) = [RLP] -> RLP
List ([RLP] -> RLP) -> [RLP] -> RLP
forall a b. (a -> b) -> a -> b
$ Seq RLP -> [RLP]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ((Ref -> RLP) -> Seq Ref -> Seq RLP
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ref -> RLP
rlpRef Seq Ref
refs) [RLP] -> [RLP] -> [RLP]
forall a. Semigroup a => a -> a -> a
<> [ByteString -> RLP
BS ByteString
val]

type NodeDB = DB ByteString Node

instance Show (NodeDB Node) where
  show :: NodeDB Node -> String
show = NodeDB Node -> String
forall a. Show a => a -> String
show

putNode :: Node -> NodeDB Ref
putNode :: Node -> NodeDB Ref
putNode Node
node =
  let bytes :: ByteString
bytes = RLP -> ByteString
rlpencode (RLP -> ByteString) -> RLP -> ByteString
forall a b. (a -> b) -> a -> b
$ Node -> RLP
rlpNode Node
node
      digest :: ByteString
digest = W256 -> ByteString
word256Bytes (W256 -> ByteString) -> W256 -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> W256
keccak ByteString
bytes
  in if ByteString -> Int
BS.length ByteString
bytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
32
    then Ref -> NodeDB Ref
forall (m :: * -> *) a. Monad m => a -> m a
return (Ref -> NodeDB Ref) -> Ref -> NodeDB Ref
forall a b. (a -> b) -> a -> b
$ Node -> Ref
Literal Node
node
    else do
      ByteString -> Node -> DB ByteString Node ()
forall k v. k -> v -> DB k v ()
insertDB ByteString
digest Node
node
      Ref -> NodeDB Ref
forall (m :: * -> *) a. Monad m => a -> m a
return (Ref -> NodeDB Ref) -> Ref -> NodeDB Ref
forall a b. (a -> b) -> a -> b
$ ByteString -> Ref
Hash ByteString
digest

getNode :: Ref -> NodeDB Node
getNode :: Ref -> NodeDB Node
getNode (Hash ByteString
d) = ByteString -> NodeDB Node
forall k v. k -> DB k v v
lookupDB ByteString
d
getNode (Literal Node
n) = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
n

lookupPath :: Ref -> Path -> NodeDB ByteString
lookupPath :: Ref -> Path -> NodeDB ByteString
lookupPath Ref
root Path
path = Ref -> NodeDB Node
getNode Ref
root NodeDB Node -> (Node -> NodeDB ByteString) -> NodeDB ByteString
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Path -> Node -> NodeDB ByteString
getVal Path
path

getVal :: Path -> Node -> NodeDB ByteString
getVal :: Path -> Node -> NodeDB ByteString
getVal Path
_ Node
Empty = ByteString -> NodeDB ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
BS.empty
getVal Path
path (Shortcut Path
nodePath Either Ref ByteString
ref) =
  case (Path -> Path -> Maybe Path
forall a. Eq a => [a] -> [a] -> Maybe [a]
stripPrefix Path
nodePath Path
path, Either Ref ByteString
ref) of
    (Just [], Right ByteString
value) -> ByteString -> NodeDB ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
value
    (Just Path
remaining, Left Ref
key) -> Ref -> Path -> NodeDB ByteString
lookupPath Ref
key Path
remaining
    (Maybe Path, Either Ref ByteString)
_ -> ByteString -> NodeDB ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
BS.empty

getVal [] (Full Seq Ref
_ ByteString
val) = ByteString -> NodeDB ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
val
getVal (Nibble
p:Path
ps) (Full Seq Ref
refs ByteString
_) = Ref -> Path -> NodeDB ByteString
lookupPath (Seq Ref
refs Seq Ref -> Int -> Ref
forall a. Seq a -> Int -> a
`Seq.index` (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
p)) Path
ps

emptyRef :: Ref
emptyRef :: Ref
emptyRef = Node -> Ref
Literal Node
Empty

emptyRefs :: Seq Ref
emptyRefs :: Seq Ref
emptyRefs = Int -> Ref -> Seq Ref
forall a. Int -> a -> Seq a
Seq.replicate Int
16 Ref
emptyRef

addPrefix :: Path -> Node -> NodeDB Node
addPrefix :: Path -> Node -> NodeDB Node
addPrefix Path
_ Node
Empty = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
Empty
addPrefix [] Node
node = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
node
addPrefix Path
path (Shortcut Path
p Either Ref ByteString
v) = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Path -> Either Ref ByteString -> Node
Shortcut (Path
path Path -> Path -> Path
forall a. Semigroup a => a -> a -> a
<> Path
p) Either Ref ByteString
v
addPrefix Path
path Node
n = Path -> Either Ref ByteString -> Node
Shortcut Path
path (Either Ref ByteString -> Node)
-> (Ref -> Either Ref ByteString) -> Ref -> Node
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ref -> Either Ref ByteString
forall a b. a -> Either a b
Left (Ref -> Node) -> NodeDB Ref -> NodeDB Node
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Node -> NodeDB Ref
putNode Node
n

insertRef :: Ref -> Path -> ByteString -> NodeDB Ref
insertRef :: Ref -> Path -> ByteString -> NodeDB Ref
insertRef Ref
ref Path
p ByteString
val = do Node
root <- Ref -> NodeDB Node
getNode Ref
ref
                         Node
newNode <- if ByteString
val ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== ByteString
BS.empty
                                    then Node -> Path -> NodeDB Node
delete Node
root Path
p
                                    else Node -> Path -> ByteString -> NodeDB Node
update Node
root Path
p ByteString
val
                         Node -> NodeDB Ref
putNode Node
newNode

update :: Node -> Path -> ByteString -> NodeDB Node
update :: Node -> Path -> ByteString -> NodeDB Node
update Node
Empty Path
p ByteString
new  = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Path -> Either Ref ByteString -> Node
Shortcut Path
p (ByteString -> Either Ref ByteString
forall a b. b -> Either a b
Right ByteString
new)
update (Full Seq Ref
refs ByteString
_) [] ByteString
new = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq Ref -> ByteString -> Node
Full Seq Ref
refs ByteString
new)
update (Full Seq Ref
refs ByteString
old) (Nibble
p:Path
ps) ByteString
new = do
  Ref
newRef <- Ref -> Path -> ByteString -> NodeDB Ref
insertRef (Seq Ref
refs Seq Ref -> Int -> Ref
forall a. Seq a -> Int -> a
`Seq.index` (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
p)) Path
ps ByteString
new
  Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Seq Ref -> ByteString -> Node
Full (Int -> Ref -> Seq Ref -> Seq Ref
forall a. Int -> a -> Seq a -> Seq a
Seq.update (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
p) Ref
newRef Seq Ref
refs) ByteString
old
update (Shortcut (Nibble
o:Path
os) (Right ByteString
old)) [] ByteString
new = do
  Ref
newRef <- Ref -> Path -> ByteString -> NodeDB Ref
insertRef Ref
emptyRef Path
os ByteString
old
  Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Seq Ref -> ByteString -> Node
Full (Int -> Ref -> Seq Ref -> Seq Ref
forall a. Int -> a -> Seq a -> Seq a
Seq.update (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
o) Ref
newRef Seq Ref
emptyRefs) ByteString
new
update (Shortcut [] (Right ByteString
old)) (Nibble
p:Path
ps) ByteString
new = do
  Ref
newRef <- Ref -> Path -> ByteString -> NodeDB Ref
insertRef Ref
emptyRef Path
ps ByteString
new
  Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Seq Ref -> ByteString -> Node
Full (Int -> Ref -> Seq Ref -> Seq Ref
forall a. Int -> a -> Seq a -> Seq a
Seq.update (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
p) Ref
newRef Seq Ref
emptyRefs) ByteString
old
update (Shortcut [] (Right ByteString
_)) [] ByteString
new =
  Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Path -> Either Ref ByteString -> Node
Shortcut [] (ByteString -> Either Ref ByteString
forall a b. b -> Either a b
Right ByteString
new)
update (Shortcut (Nibble
o:Path
os) Either Ref ByteString
to) (Nibble
p:Path
ps) ByteString
new | Nibble
o Nibble -> Nibble -> Bool
forall a. Eq a => a -> a -> Bool
== Nibble
p
  = Node -> Path -> ByteString -> NodeDB Node
update (Path -> Either Ref ByteString -> Node
Shortcut Path
os Either Ref ByteString
to) Path
ps ByteString
new NodeDB Node -> (Node -> NodeDB Node) -> NodeDB Node
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Path -> Node -> NodeDB Node
addPrefix [Nibble
o]
                                       | Bool
otherwise = do
  Ref
oldRef <- case Either Ref ByteString
to of
              (Left Ref
ref)  -> Ref -> NodeDB Node
getNode Ref
ref NodeDB Node -> (Node -> NodeDB Node) -> NodeDB Node
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Path -> Node -> NodeDB Node
addPrefix Path
os NodeDB Node -> (Node -> NodeDB Ref) -> NodeDB Ref
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Node -> NodeDB Ref
putNode
              (Right ByteString
val) -> Ref -> Path -> ByteString -> NodeDB Ref
insertRef Ref
emptyRef Path
os ByteString
val
  Ref
newRef <- Ref -> Path -> ByteString -> NodeDB Ref
insertRef Ref
emptyRef Path
ps ByteString
new
  let refs :: Seq Ref
refs = Int -> Ref -> Seq Ref -> Seq Ref
forall a. Int -> a -> Seq a -> Seq a
Seq.update (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
p) Ref
newRef (Seq Ref -> Seq Ref) -> Seq Ref -> Seq Ref
forall a b. (a -> b) -> a -> b
$ Int -> Ref -> Seq Ref -> Seq Ref
forall a. Int -> a -> Seq a -> Seq a
Seq.update (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
o) Ref
oldRef Seq Ref
emptyRefs
  Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Seq Ref -> ByteString -> Node
Full Seq Ref
refs ByteString
BS.empty
update (Shortcut (Nibble
o:Path
os) (Left Ref
ref)) [] ByteString
new = do
  Ref
newRef <- Ref -> NodeDB Node
getNode Ref
ref NodeDB Node -> (Node -> NodeDB Node) -> NodeDB Node
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Path -> Node -> NodeDB Node
addPrefix Path
os NodeDB Node -> (Node -> NodeDB Ref) -> NodeDB Ref
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Node -> NodeDB Ref
putNode
  Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Seq Ref -> ByteString -> Node
Full (Int -> Ref -> Seq Ref -> Seq Ref
forall a. Int -> a -> Seq a -> Seq a
Seq.update (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
o) Ref
newRef Seq Ref
emptyRefs) ByteString
new
update (Shortcut Path
cut (Left Ref
ref)) Path
ps ByteString
new = do
  Ref
newRef <- Ref -> Path -> ByteString -> NodeDB Ref
insertRef Ref
ref Path
ps ByteString
new
  Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Path -> Either Ref ByteString -> Node
Shortcut Path
cut (Ref -> Either Ref ByteString
forall a b. a -> Either a b
Left Ref
newRef)

delete :: Node -> Path -> NodeDB Node
delete :: Node -> Path -> NodeDB Node
delete Node
Empty Path
_ = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
Empty
delete (Shortcut [] (Right ByteString
_)) [] = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
Empty
delete n :: Node
n@(Shortcut [] (Right ByteString
_)) Path
_ = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
n
delete (Shortcut [] (Left Ref
ref)) Path
p = do Node
node <- Ref -> NodeDB Node
getNode Ref
ref
                                       Node -> Path -> NodeDB Node
delete Node
node Path
p
delete n :: Node
n@(Shortcut Path
_ Either Ref ByteString
_) [] = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
n
delete n :: Node
n@(Shortcut (Nibble
o:Path
os) Either Ref ByteString
to) (Nibble
p:Path
ps) | Nibble
p Nibble -> Nibble -> Bool
forall a. Eq a => a -> a -> Bool
== Nibble
o
  = Node -> Path -> NodeDB Node
delete (Path -> Either Ref ByteString -> Node
Shortcut Path
os Either Ref ByteString
to) Path
ps NodeDB Node -> (Node -> NodeDB Node) -> NodeDB Node
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Path -> Node -> NodeDB Node
addPrefix [Nibble
o]
                                     | Bool
otherwise
  = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
n
delete (Full Seq Ref
refs ByteString
_) [] | Seq Ref
refs Seq Ref -> Seq Ref -> Bool
forall a. Eq a => a -> a -> Bool
== Seq Ref
emptyRefs
  = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
Empty
                        | Bool
otherwise
  = Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq Ref -> ByteString -> Node
Full Seq Ref
refs ByteString
BS.empty)
delete (Full Seq Ref
refs ByteString
val) (Nibble
p:Path
ps) = do
  Ref
newRef <- Ref -> Path -> ByteString -> NodeDB Ref
insertRef (Seq Ref
refs Seq Ref -> Int -> Ref
forall a. Seq a -> Int -> a
`Seq.index` (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
p)) Path
ps ByteString
BS.empty
  let newRefs :: Seq Ref
newRefs = Int -> Ref -> Seq Ref -> Seq Ref
forall a. Int -> a -> Seq a -> Seq a
Seq.update (Nibble -> Int
forall a b. (Integral a, Num b) => a -> b
num Nibble
p) Ref
newRef Seq Ref
refs
      nonEmpties :: [(Word8, Ref)]
nonEmpties = ((Word8, Ref) -> Bool) -> [(Word8, Ref)] -> [(Word8, Ref)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(Word8
_, Ref
ref) -> Ref
ref Ref -> Ref -> Bool
forall a. Eq a => a -> a -> Bool
/= Ref
emptyRef) ([(Word8, Ref)] -> [(Word8, Ref)])
-> [(Word8, Ref)] -> [(Word8, Ref)]
forall a b. (a -> b) -> a -> b
$ [Word8] -> [Ref] -> [(Word8, Ref)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Word8
0..Word8
15] ([Ref] -> [(Word8, Ref)]) -> [Ref] -> [(Word8, Ref)]
forall a b. (a -> b) -> a -> b
$ Seq Ref -> [Ref]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Seq Ref
newRefs
  case ([(Word8, Ref)]
nonEmpties, ByteString -> Bool
BS.null ByteString
val) of
    ([], Bool
True)         -> Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return Node
Empty
    ([(Word8
n, Ref
ref)], Bool
True)  -> Ref -> NodeDB Node
getNode Ref
ref NodeDB Node -> (Node -> NodeDB Node) -> NodeDB Node
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Path -> Node -> NodeDB Node
addPrefix [Word8 -> Nibble
Nibble Word8
n]
    ([(Word8, Ref)], Bool)
_                    -> Node -> NodeDB Node
forall (m :: * -> *) a. Monad m => a -> m a
return (Node -> NodeDB Node) -> Node -> NodeDB Node
forall a b. (a -> b) -> a -> b
$ Seq Ref -> ByteString -> Node
Full Seq Ref
newRefs ByteString
val

insert :: Ref -> ByteString -> ByteString -> NodeDB Ref
insert :: Ref -> ByteString -> ByteString -> NodeDB Ref
insert Ref
ref ByteString
key = Ref -> Path -> ByteString -> NodeDB Ref
insertRef Ref
ref (ByteString -> Path
unpackNibbles ByteString
key)

lookupIn :: Ref -> ByteString -> NodeDB ByteString
lookupIn :: Ref -> ByteString -> NodeDB ByteString
lookupIn Ref
ref ByteString
bs = Ref -> Path -> NodeDB ByteString
lookupPath Ref
ref (Path -> NodeDB ByteString) -> Path -> NodeDB ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> Path
unpackNibbles ByteString
bs

type Trie = StateT Ref NodeDB

runTrie :: DB ByteString ByteString a -> Trie a
runTrie :: DB ByteString ByteString a -> Trie a
runTrie = (ByteString -> ByteString -> StateT Ref (DB ByteString Node) ())
-> (ByteString -> StateT Ref (DB ByteString Node) ByteString)
-> DB ByteString ByteString a
-> Trie a
forall (m :: * -> *) k v a.
Monad m =>
(k -> v -> m ()) -> (k -> m v) -> DB k v a -> m a
runDB ByteString -> ByteString -> StateT Ref (DB ByteString Node) ()
forall (t :: (* -> *) -> * -> *).
(MonadState Ref (t (DB ByteString Node)), MonadTrans t) =>
ByteString -> ByteString -> t (DB ByteString Node) ()
putDB ByteString -> StateT Ref (DB ByteString Node) ByteString
forall (t :: (* -> *) -> * -> *).
(MonadState Ref (t (DB ByteString Node)), MonadTrans t) =>
ByteString -> t (DB ByteString Node) ByteString
getDB
  where
    putDB :: ByteString -> ByteString -> t (DB ByteString Node) ()
putDB ByteString
key ByteString
val = do
      Ref
ref <- t (DB ByteString Node) Ref
forall s (m :: * -> *). MonadState s m => m s
get
      Ref
newRef <- NodeDB Ref -> t (DB ByteString Node) Ref
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (NodeDB Ref -> t (DB ByteString Node) Ref)
-> NodeDB Ref -> t (DB ByteString Node) Ref
forall a b. (a -> b) -> a -> b
$ Ref -> ByteString -> ByteString -> NodeDB Ref
insert Ref
ref ByteString
key ByteString
val
      Ref -> t (DB ByteString Node) ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put Ref
newRef
    getDB :: ByteString -> t (DB ByteString Node) ByteString
getDB ByteString
key = do
      Ref
ref <- t (DB ByteString Node) Ref
forall s (m :: * -> *). MonadState s m => m s
get
      NodeDB ByteString -> t (DB ByteString Node) ByteString
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (NodeDB ByteString -> t (DB ByteString Node) ByteString)
-> NodeDB ByteString -> t (DB ByteString Node) ByteString
forall a b. (a -> b) -> a -> b
$ Ref -> ByteString -> NodeDB ByteString
lookupIn Ref
ref ByteString
key

type MapDB k v a = StateT (Map.Map k v) Maybe a

runMapDB :: Ord k => DB k v a -> MapDB k v a
runMapDB :: DB k v a -> MapDB k v a
runMapDB = (k -> v -> StateT (Map k v) Maybe ())
-> (k -> StateT (Map k v) Maybe v) -> DB k v a -> MapDB k v a
forall (m :: * -> *) k v a.
Monad m =>
(k -> v -> m ()) -> (k -> m v) -> DB k v a -> m a
runDB k -> v -> StateT (Map k v) Maybe ()
forall (m :: * -> *) k a.
(MonadState (Map k a) m, Ord k) =>
k -> a -> m ()
putDB k -> StateT (Map k v) Maybe v
forall (t :: (* -> *) -> * -> *) k b.
(MonadState (Map k b) (t Maybe), MonadTrans t, Ord k) =>
k -> t Maybe b
getDB
  where
    getDB :: k -> t Maybe b
getDB k
key = do
      Map k b
mmap <- t Maybe (Map k b)
forall s (m :: * -> *). MonadState s m => m s
get
      Maybe b -> t Maybe b
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Maybe b -> t Maybe b) -> Maybe b -> t Maybe b
forall a b. (a -> b) -> a -> b
$ k -> Map k b -> Maybe b
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
key Map k b
mmap
    putDB :: k -> a -> m ()
putDB k
key a
value = do
      Map k a
mmap <- m (Map k a)
forall s (m :: * -> *). MonadState s m => m s
get
      let newMap :: Map k a
newMap = k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
key a
value Map k a
mmap
      Map k a -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put Map k a
newMap


insertValues :: [(ByteString, ByteString)] -> Maybe Ref
insertValues :: [(ByteString, ByteString)] -> Maybe Ref
insertValues [(ByteString, ByteString)]
inputs =
  let trie :: StateT Ref (DB ByteString Node) ()
trie = DB ByteString ByteString () -> StateT Ref (DB ByteString Node) ()
forall a. DB ByteString ByteString a -> Trie a
runTrie (DB ByteString ByteString () -> StateT Ref (DB ByteString Node) ())
-> DB ByteString ByteString ()
-> StateT Ref (DB ByteString Node) ()
forall a b. (a -> b) -> a -> b
$ ((ByteString, ByteString) -> DB ByteString ByteString ())
-> [(ByteString, ByteString)] -> DB ByteString ByteString ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (ByteString, ByteString) -> DB ByteString ByteString ()
forall k v. (k, v) -> DB k v ()
insertPair [(ByteString, ByteString)]
inputs
      mapDB :: MapDB ByteString Node ((), Ref)
mapDB = DB ByteString Node ((), Ref) -> MapDB ByteString Node ((), Ref)
forall k v a. Ord k => DB k v a -> MapDB k v a
runMapDB (DB ByteString Node ((), Ref) -> MapDB ByteString Node ((), Ref))
-> DB ByteString Node ((), Ref) -> MapDB ByteString Node ((), Ref)
forall a b. (a -> b) -> a -> b
$ StateT Ref (DB ByteString Node) ()
-> Ref -> DB ByteString Node ((), Ref)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT StateT Ref (DB ByteString Node) ()
trie (Node -> Ref
Literal Node
Empty)
      result :: Maybe Ref
result = ((), Ref) -> Ref
forall a b. (a, b) -> b
snd (((), Ref) -> Ref) -> Maybe ((), Ref) -> Maybe Ref
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MapDB ByteString Node ((), Ref)
-> Map ByteString Node -> Maybe ((), Ref)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT MapDB ByteString Node ((), Ref)
mapDB Map ByteString Node
forall k a. Map k a
Map.empty
      insertPair :: (k, v) -> DB k v ()
insertPair (k
key, v
value) = k -> v -> DB k v ()
forall k v. k -> v -> DB k v ()
insertDB k
key v
value
  in Maybe Ref
result

calcRoot :: [(ByteString, ByteString)] -> Maybe ByteString
calcRoot :: [(ByteString, ByteString)] -> Maybe ByteString
calcRoot [(ByteString, ByteString)]
vs = case [(ByteString, ByteString)] -> Maybe Ref
insertValues [(ByteString, ByteString)]
vs of
     Just (Hash ByteString
b) -> ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
b
     Just (Literal Node
n) -> ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString) -> ByteString -> Maybe ByteString
forall a b. (a -> b) -> a -> b
$ W256 -> ByteString
word256Bytes (W256 -> ByteString) -> W256 -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> W256
keccak (ByteString -> W256) -> ByteString -> W256
forall a b. (a -> b) -> a -> b
$ RLP -> ByteString
rlpencode (RLP -> ByteString) -> RLP -> ByteString
forall a b. (a -> b) -> a -> b
$ Node -> RLP
rlpNode Node
n
     Maybe Ref
Nothing -> Maybe ByteString
forall a. Maybe a
Nothing