From ca7acf6b6373182ffd4681587f3becb15b8c9c80 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Fri, 15 Nov 2024 13:10:59 +0200
Subject: [PATCH 01/31] feat: haskell project setup

---
 .gitignore                    |  1 +
 pubsub-haskell/.envrc         |  1 +
 pubsub-haskell/CHANGELOG.md   |  5 +++
 pubsub-haskell/app/Main.hs    |  4 ++
 pubsub-haskell/flake.lock     | 61 +++++++++++++++++++++++++++
 pubsub-haskell/flake.nix      | 30 ++++++++++++++
 pubsub-haskell/sky-node.cabal | 77 +++++++++++++++++++++++++++++++++++
 7 files changed, 179 insertions(+)
 create mode 100644 pubsub-haskell/.envrc
 create mode 100644 pubsub-haskell/CHANGELOG.md
 create mode 100644 pubsub-haskell/app/Main.hs
 create mode 100644 pubsub-haskell/flake.lock
 create mode 100644 pubsub-haskell/flake.nix
 create mode 100644 pubsub-haskell/sky-node.cabal

diff --git a/.gitignore b/.gitignore
index 5900d22..663e461 100644
--- a/.gitignore
+++ b/.gitignore
@@ -36,3 +36,4 @@ TAGS
 
 # Cache and other dev setup files
 .direnv/
+*/dist-newstyle/*
diff --git a/pubsub-haskell/.envrc b/pubsub-haskell/.envrc
new file mode 100644
index 0000000..3550a30
--- /dev/null
+++ b/pubsub-haskell/.envrc
@@ -0,0 +1 @@
+use flake
diff --git a/pubsub-haskell/CHANGELOG.md b/pubsub-haskell/CHANGELOG.md
new file mode 100644
index 0000000..346c814
--- /dev/null
+++ b/pubsub-haskell/CHANGELOG.md
@@ -0,0 +1,5 @@
+# Revision history for sky-node
+
+## 0.1.0.0 -- YYYY-mm-dd
+
+* First version. Released on an unsuspecting world.
diff --git a/pubsub-haskell/app/Main.hs b/pubsub-haskell/app/Main.hs
new file mode 100644
index 0000000..65ae4a0
--- /dev/null
+++ b/pubsub-haskell/app/Main.hs
@@ -0,0 +1,4 @@
+module Main where
+
+main :: IO ()
+main = putStrLn "Hello, Haskell!"
diff --git a/pubsub-haskell/flake.lock b/pubsub-haskell/flake.lock
new file mode 100644
index 0000000..d7d724e
--- /dev/null
+++ b/pubsub-haskell/flake.lock
@@ -0,0 +1,61 @@
+{
+  "nodes": {
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1731319897,
+        "narHash": "sha256-PbABj4tnbWFMfBp6OcUK5iGy1QY+/Z96ZcLpooIbuEI=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "dc460ec76cbff0e66e269457d7b728432263166c",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}
diff --git a/pubsub-haskell/flake.nix b/pubsub-haskell/flake.nix
new file mode 100644
index 0000000..b08fd04
--- /dev/null
+++ b/pubsub-haskell/flake.nix
@@ -0,0 +1,30 @@
+{
+  inputs = {
+    flake-utils.url = "github:numtide/flake-utils";
+    nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
+  };
+  outputs = { self, nixpkgs, flake-utils, ... }: flake-utils.lib.eachDefaultSystem (system:
+    let
+      pkgs = import nixpkgs { inherit system; };
+      overlay = final: prev: {
+        sky-node = final.callCabal2nix "sky-node" ./. { };
+      };
+      haskPkgs = pkgs.haskellPackages.extend overlay;
+    in
+    {
+      devShells.default = haskPkgs.shellFor {
+        packages = p: [
+          p.sky-node
+        ];
+        nativeBuildInputs = with haskPkgs; [
+          cabal-install
+          haskell-language-server
+          hlint
+        ];
+        shellHook = ''
+          echo "Welcome"
+        '';
+      };
+    }
+  );
+}
diff --git a/pubsub-haskell/sky-node.cabal b/pubsub-haskell/sky-node.cabal
new file mode 100644
index 0000000..822857a
--- /dev/null
+++ b/pubsub-haskell/sky-node.cabal
@@ -0,0 +1,77 @@
+cabal-version:      3.0
+-- The cabal-version field refers to the version of the .cabal specification,
+-- and can be different from the cabal-install (the tool) version and the
+-- Cabal (the library) version you are using. As such, the Cabal (the library)
+-- version used must be equal or greater than the version stated in this field.
+-- Starting from the specification version 2.2, the cabal-version field must be
+-- the first thing in the cabal file.
+
+-- Initial package description 'sky-node' generated by
+-- 'cabal init'. For further documentation, see:
+--   http://haskell.org/cabal/users-guide/
+--
+-- The name of the package.
+name:               sky-node
+
+-- The package version.
+-- See the Haskell package versioning policy (PVP) for standards
+-- guiding when and how versions should be incremented.
+-- https://pvp.haskell.org
+-- PVP summary:     +-+------- breaking API changes
+--                  | | +----- non-breaking API additions
+--                  | | | +--- code changes with no API change
+version:            0.1.0.0
+
+-- A short (one-line) description of the package.
+synopsis:           Data Availability network for Cardano
+
+-- A longer description of the package.
+-- description:
+
+-- URL for the project homepage or repository.
+homepage:           https://www.skyprotocol.org/
+
+-- The license under which the package is released.
+license:            NONE
+
+-- The package author(s).
+author:             Yaroslav Kozhevnikov
+
+-- An email address to which users can send suggestions, bug reports, and patches.
+maintainer:         yroslav541@gmail.com
+
+-- A copyright notice.
+-- copyright:
+category:           Network
+build-type:         Simple
+
+-- Extra doc files to be distributed with the package, such as a CHANGELOG or a README.
+extra-doc-files:    CHANGELOG.md
+
+-- Extra source files to be distributed with the package, such as examples, or a tutorial module.
+-- extra-source-files:
+
+common warnings
+    ghc-options: -Wall
+
+executable sky-node
+    -- Import common warning flags.
+    import:           warnings
+
+    -- .hs or .lhs file containing the Main module.
+    main-is:          Main.hs
+
+    -- Modules included in this executable, other than Main.
+    -- other-modules:
+
+    -- LANGUAGE extensions used by modules in this package.
+    -- other-extensions:
+
+    -- Other library packages from which modules are imported.
+    build-depends:    base ^>=4.18.2.1
+
+    -- Directories containing source files.
+    hs-source-dirs:   app
+
+    -- Base language which the package is written in.
+    default-language: Haskell2010

From e5c8deb3ecf2da761dc9541cef572e0929442cf2 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Fri, 15 Nov 2024 15:43:09 +0200
Subject: [PATCH 02/31] feat: base project structure & more setup

---
 pubsub-haskell/.envrc              |  1 +
 pubsub-haskell/app/Main.hs         | 11 ++++++++++-
 pubsub-haskell/config/default.yaml |  1 +
 pubsub-haskell/hie.yaml            |  7 +++++++
 pubsub-haskell/lib/App.hs          | 16 ++++++++++++++++
 pubsub-haskell/lib/App/Error.hs    |  6 ++++++
 pubsub-haskell/lib/Config.hs       | 12 ++++++++++++
 pubsub-haskell/sky-node.cabal      | 30 +++++++++++++++++++++++++++---
 8 files changed, 80 insertions(+), 4 deletions(-)
 create mode 100644 pubsub-haskell/config/default.yaml
 create mode 100644 pubsub-haskell/hie.yaml
 create mode 100644 pubsub-haskell/lib/App.hs
 create mode 100644 pubsub-haskell/lib/App/Error.hs
 create mode 100644 pubsub-haskell/lib/Config.hs

diff --git a/pubsub-haskell/.envrc b/pubsub-haskell/.envrc
index 3550a30..a971e5d 100644
--- a/pubsub-haskell/.envrc
+++ b/pubsub-haskell/.envrc
@@ -1 +1,2 @@
+watch_file *.cabal
 use flake
diff --git a/pubsub-haskell/app/Main.hs b/pubsub-haskell/app/Main.hs
index 65ae4a0..d5e2e31 100644
--- a/pubsub-haskell/app/Main.hs
+++ b/pubsub-haskell/app/Main.hs
@@ -1,4 +1,13 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Main where
 
+import App (runApp)
+import Data.Yaml.Config (loadYamlSettingsArgs, useEnv)
+import System.Log.FastLogger (defaultBufSize, newStdoutLoggerSet)
+
 main :: IO ()
-main = putStrLn "Hello, Haskell!"
+main = do
+  logger <- newStdoutLoggerSet defaultBufSize
+  config <- loadYamlSettingsArgs ["config/default.yaml"] useEnv
+  runApp config logger
diff --git a/pubsub-haskell/config/default.yaml b/pubsub-haskell/config/default.yaml
new file mode 100644
index 0000000..d130112
--- /dev/null
+++ b/pubsub-haskell/config/default.yaml
@@ -0,0 +1 @@
+port: "3000"
diff --git a/pubsub-haskell/hie.yaml b/pubsub-haskell/hie.yaml
new file mode 100644
index 0000000..230dbb8
--- /dev/null
+++ b/pubsub-haskell/hie.yaml
@@ -0,0 +1,7 @@
+cradle:
+  cabal:
+    - path: "lib"
+      component: "sky-node:lib:sky-lib"
+
+    - path: "app/Main.hs"
+      component: "sky-node:exe:sky-node"
diff --git a/pubsub-haskell/lib/App.hs b/pubsub-haskell/lib/App.hs
new file mode 100644
index 0000000..99e759e
--- /dev/null
+++ b/pubsub-haskell/lib/App.hs
@@ -0,0 +1,16 @@
+module App (runApp) where
+
+-- import App.Error
+import Config (AppConfig)
+import System.Log.FastLogger (LoggerSet, pushLogStrLn, toLogStr)
+
+-- initAppState :: IO ()
+-- initAppState = pure ()
+
+runApp :: AppConfig -> LoggerSet -> IO ()
+runApp config logger = do
+  -- state <- initAppState
+  print config
+  pushLogStrLn logger $ toLogStr "Starting Sky Node"
+
+-- runNode config state logger
diff --git a/pubsub-haskell/lib/App/Error.hs b/pubsub-haskell/lib/App/Error.hs
new file mode 100644
index 0000000..3f778f7
--- /dev/null
+++ b/pubsub-haskell/lib/App/Error.hs
@@ -0,0 +1,6 @@
+module App.Error where
+
+data AppError
+  = HandlerError String
+  | AppError String
+  deriving (Show)
diff --git a/pubsub-haskell/lib/Config.hs b/pubsub-haskell/lib/Config.hs
new file mode 100644
index 0000000..c5ba9c6
--- /dev/null
+++ b/pubsub-haskell/lib/Config.hs
@@ -0,0 +1,12 @@
+{-# LANGUAGE DeriveAnyClass #-}
+{-# LANGUAGE DeriveGeneric #-}
+
+module Config where
+
+import Data.Yaml (FromJSON)
+import GHC.Generics (Generic)
+
+data AppConfig = AppConfig
+  { port :: String
+  }
+  deriving (Show, Generic, FromJSON)
diff --git a/pubsub-haskell/sky-node.cabal b/pubsub-haskell/sky-node.cabal
index 822857a..c3286eb 100644
--- a/pubsub-haskell/sky-node.cabal
+++ b/pubsub-haskell/sky-node.cabal
@@ -54,9 +54,32 @@ extra-doc-files:    CHANGELOG.md
 common warnings
     ghc-options: -Wall
 
+common shared-libs
+    build-depends:
+        base ^>=4.18.2.1,
+        yaml,
+        fast-logger
+
+library sky-lib
+    import: shared-libs, warnings
+
+    exposed-modules:
+        App
+        App.Error
+        Config
+
+    build-depends:
+        network,
+        mtl,
+        exceptions,
+
+    hs-source-dirs: lib
+
+    default-language: Haskell2010
+
 executable sky-node
     -- Import common warning flags.
-    import:           warnings
+    import: shared-libs, warnings
 
     -- .hs or .lhs file containing the Main module.
     main-is:          Main.hs
@@ -68,10 +91,11 @@ executable sky-node
     -- other-extensions:
 
     -- Other library packages from which modules are imported.
-    build-depends:    base ^>=4.18.2.1
+    build-depends:
+        sky-lib
 
     -- Directories containing source files.
-    hs-source-dirs:   app
+    hs-source-dirs: app
 
     -- Base language which the package is written in.
     default-language: Haskell2010

From 93c445b25a6a209e896442bfdf94a458d9554685 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Fri, 15 Nov 2024 18:07:01 +0200
Subject: [PATCH 03/31] chore: rename

---
 {pubsub-haskell => sky-node}/.envrc              | 0
 {pubsub-haskell => sky-node}/CHANGELOG.md        | 0
 {pubsub-haskell => sky-node}/app/Main.hs         | 0
 {pubsub-haskell => sky-node}/config/default.yaml | 0
 {pubsub-haskell => sky-node}/flake.lock          | 0
 {pubsub-haskell => sky-node}/flake.nix           | 0
 {pubsub-haskell => sky-node}/hie.yaml            | 0
 {pubsub-haskell => sky-node}/lib/App.hs          | 0
 {pubsub-haskell => sky-node}/lib/App/Error.hs    | 0
 {pubsub-haskell => sky-node}/lib/Config.hs       | 0
 {pubsub-haskell => sky-node}/sky-node.cabal      | 0
 11 files changed, 0 insertions(+), 0 deletions(-)
 rename {pubsub-haskell => sky-node}/.envrc (100%)
 rename {pubsub-haskell => sky-node}/CHANGELOG.md (100%)
 rename {pubsub-haskell => sky-node}/app/Main.hs (100%)
 rename {pubsub-haskell => sky-node}/config/default.yaml (100%)
 rename {pubsub-haskell => sky-node}/flake.lock (100%)
 rename {pubsub-haskell => sky-node}/flake.nix (100%)
 rename {pubsub-haskell => sky-node}/hie.yaml (100%)
 rename {pubsub-haskell => sky-node}/lib/App.hs (100%)
 rename {pubsub-haskell => sky-node}/lib/App/Error.hs (100%)
 rename {pubsub-haskell => sky-node}/lib/Config.hs (100%)
 rename {pubsub-haskell => sky-node}/sky-node.cabal (100%)

diff --git a/pubsub-haskell/.envrc b/sky-node/.envrc
similarity index 100%
rename from pubsub-haskell/.envrc
rename to sky-node/.envrc
diff --git a/pubsub-haskell/CHANGELOG.md b/sky-node/CHANGELOG.md
similarity index 100%
rename from pubsub-haskell/CHANGELOG.md
rename to sky-node/CHANGELOG.md
diff --git a/pubsub-haskell/app/Main.hs b/sky-node/app/Main.hs
similarity index 100%
rename from pubsub-haskell/app/Main.hs
rename to sky-node/app/Main.hs
diff --git a/pubsub-haskell/config/default.yaml b/sky-node/config/default.yaml
similarity index 100%
rename from pubsub-haskell/config/default.yaml
rename to sky-node/config/default.yaml
diff --git a/pubsub-haskell/flake.lock b/sky-node/flake.lock
similarity index 100%
rename from pubsub-haskell/flake.lock
rename to sky-node/flake.lock
diff --git a/pubsub-haskell/flake.nix b/sky-node/flake.nix
similarity index 100%
rename from pubsub-haskell/flake.nix
rename to sky-node/flake.nix
diff --git a/pubsub-haskell/hie.yaml b/sky-node/hie.yaml
similarity index 100%
rename from pubsub-haskell/hie.yaml
rename to sky-node/hie.yaml
diff --git a/pubsub-haskell/lib/App.hs b/sky-node/lib/App.hs
similarity index 100%
rename from pubsub-haskell/lib/App.hs
rename to sky-node/lib/App.hs
diff --git a/pubsub-haskell/lib/App/Error.hs b/sky-node/lib/App/Error.hs
similarity index 100%
rename from pubsub-haskell/lib/App/Error.hs
rename to sky-node/lib/App/Error.hs
diff --git a/pubsub-haskell/lib/Config.hs b/sky-node/lib/Config.hs
similarity index 100%
rename from pubsub-haskell/lib/Config.hs
rename to sky-node/lib/Config.hs
diff --git a/pubsub-haskell/sky-node.cabal b/sky-node/sky-node.cabal
similarity index 100%
rename from pubsub-haskell/sky-node.cabal
rename to sky-node/sky-node.cabal

From a39ab705d55ecd59959685bed8884402b763b111 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Fri, 15 Nov 2024 20:00:38 +0200
Subject: [PATCH 04/31] feat: basic app structure

---
 sky-node/lib/App.hs       | 46 +++++++++++++++++++++++++++++++++------
 sky-node/lib/App/Error.hs |  7 ++++--
 sky-node/lib/Config.hs    |  1 +
 sky-node/sky-node.cabal   |  1 +
 4 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/sky-node/lib/App.hs b/sky-node/lib/App.hs
index 99e759e..d555e0f 100644
--- a/sky-node/lib/App.hs
+++ b/sky-node/lib/App.hs
@@ -1,16 +1,48 @@
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE RecordWildCards #-}
+
 module App (runApp) where
 
 -- import App.Error
-import Config (AppConfig)
+import Config (AppConfig (..))
+import Control.Concurrent.STM (newTVarIO)
+import Control.Concurrent.STM.TVar (TVar)
 import System.Log.FastLogger (LoggerSet, pushLogStrLn, toLogStr)
 
--- initAppState :: IO ()
--- initAppState = pure ()
+-- | Aux data structure to keep info about 'Peer'.
+newtype Peer = Peer
+  { -- | Id of the 'Peer'
+    id :: String
+  }
+  deriving (Show)
+
+-- | Signalizes if the node needs to shutdown or continue running.
+data Shutdown = Continue | Shutdown deriving (Show)
+
+-- | State of the node.
+data AppState = AppState
+  { -- | Messages.
+    messages :: TVar [Int],
+    -- | List of peers that are connected to this node.
+    peers :: [Peer],
+    -- | Marker for graceful shutdown.
+    continue :: Shutdown
+  }
+
+-- | Initializes default node state
+initAppState :: IO AppState
+initAppState = do
+  let peers = []
+      continue = Continue
+  messages <- newTVarIO []
+  pure $ AppState {..}
 
+-- | Run the application
 runApp :: AppConfig -> LoggerSet -> IO ()
 runApp config logger = do
-  -- state <- initAppState
-  print config
-  pushLogStrLn logger $ toLogStr "Starting Sky Node"
+  state <- initAppState
+  pushLogStrLn logger . toLogStr $ "Starting Sky Node on port " <> config.port
+  runServer config state logger
 
--- runNode config state logger
+runServer :: AppConfig -> AppState -> LoggerSet -> IO ()
+runServer config state logger = undefined
diff --git a/sky-node/lib/App/Error.hs b/sky-node/lib/App/Error.hs
index 3f778f7..cfa76b9 100644
--- a/sky-node/lib/App/Error.hs
+++ b/sky-node/lib/App/Error.hs
@@ -1,6 +1,9 @@
 module App.Error where
 
+-- | Top-level error type
 data AppError
-  = HandlerError String
-  | AppError String
+  = -- | Constructor for error propagation from the handler level. Used for errors that must terminate the execution of the node.
+    HandlerError String
+  | -- | Constructor for errors in the node itself.
+    AppError String
   deriving (Show)
diff --git a/sky-node/lib/Config.hs b/sky-node/lib/Config.hs
index c5ba9c6..aac2f6d 100644
--- a/sky-node/lib/Config.hs
+++ b/sky-node/lib/Config.hs
@@ -6,6 +6,7 @@ module Config where
 import Data.Yaml (FromJSON)
 import GHC.Generics (Generic)
 
+-- | Node configuration
 data AppConfig = AppConfig
   { port :: String
   }
diff --git a/sky-node/sky-node.cabal b/sky-node/sky-node.cabal
index c3286eb..f1c0e4e 100644
--- a/sky-node/sky-node.cabal
+++ b/sky-node/sky-node.cabal
@@ -72,6 +72,7 @@ library sky-lib
         network,
         mtl,
         exceptions,
+        stm,
 
     hs-source-dirs: lib
 

From 1c2619925773ff4b0d3780b6ddc2ed0c3cef9c04 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Sat, 16 Nov 2024 20:38:43 +0200
Subject: [PATCH 05/31] feat: runServer impl; some docs

---
 sky-node/config/default.yaml |  3 ++-
 sky-node/lib/App.hs          | 38 +++++++++++++++++++++++++++++-------
 sky-node/lib/Command.hs      |  0
 sky-node/lib/Config.hs       |  3 ++-
 sky-node/lib/Types.hs        |  0
 sky-node/sky-node.cabal      |  1 +
 6 files changed, 36 insertions(+), 9 deletions(-)
 create mode 100644 sky-node/lib/Command.hs
 create mode 100644 sky-node/lib/Types.hs

diff --git a/sky-node/config/default.yaml b/sky-node/config/default.yaml
index d130112..b1df009 100644
--- a/sky-node/config/default.yaml
+++ b/sky-node/config/default.yaml
@@ -1 +1,2 @@
-port: "3000"
+port: "3456"
+address: "localhost"
diff --git a/sky-node/lib/App.hs b/sky-node/lib/App.hs
index d555e0f..8b14885 100644
--- a/sky-node/lib/App.hs
+++ b/sky-node/lib/App.hs
@@ -7,6 +7,9 @@ module App (runApp) where
 import Config (AppConfig (..))
 import Control.Concurrent.STM (newTVarIO)
 import Control.Concurrent.STM.TVar (TVar)
+import Control.Exception (bracket)
+import Data.Default
+import qualified Network.Socket as S
 import System.Log.FastLogger (LoggerSet, pushLogStrLn, toLogStr)
 
 -- | Aux data structure to keep info about 'Peer'.
@@ -19,6 +22,9 @@ newtype Peer = Peer
 -- | Signalizes if the node needs to shutdown or continue running.
 data Shutdown = Continue | Shutdown deriving (Show)
 
+instance Default Shutdown where
+  def = Continue
+
 -- | State of the node.
 data AppState = AppState
   { -- | Messages.
@@ -29,20 +35,38 @@ data AppState = AppState
     continue :: Shutdown
   }
 
--- | Initializes default node state
+-- | Initializes default node state.
 initAppState :: IO AppState
 initAppState = do
-  let peers = []
-      continue = Continue
-  messages <- newTVarIO []
+  let peers = def
+      continue = def
+  messages <- newTVarIO def
   pure $ AppState {..}
 
--- | Run the application
+-- | Runs the application with the default 'AppState'.
 runApp :: AppConfig -> LoggerSet -> IO ()
 runApp config logger = do
   state <- initAppState
-  pushLogStrLn logger . toLogStr $ "Starting Sky Node on port " <> config.port
+  pushLogStrLn logger $ toLogStr "Starting Sky Node..."
   runServer config state logger
 
+-- | Runs the server. Binds socket to the address and accepts incoming connection.
 runServer :: AppConfig -> AppState -> LoggerSet -> IO ()
-runServer config state logger = undefined
+runServer config state logger = do
+  addr <- resolve config.hostname config.port
+  bracket (S.openSocket addr) S.close $ \sock -> do
+    S.setSocketOption sock S.ReuseAddr 1 -- easier for debugging
+    S.bind sock (S.addrAddress addr)
+    S.listen sock 10
+    pushLogStrLn logger . toLogStr $ "Node listening on port " <> config.port
+    acceptLoop sock state logger
+  where
+    resolve :: String -> String -> IO S.AddrInfo
+    resolve host port = do
+      let hints = S.defaultHints {S.addrSocketType = S.Stream}
+      addr : _ <- S.getAddrInfo (Just hints) (Just $ host <> ":" <> port) Nothing
+      pure addr
+
+-- | Loop handling new connections.
+acceptLoop :: S.Socket -> AppState -> LoggerSet -> IO ()
+acceptLoop = undefined
diff --git a/sky-node/lib/Command.hs b/sky-node/lib/Command.hs
new file mode 100644
index 0000000..e69de29
diff --git a/sky-node/lib/Config.hs b/sky-node/lib/Config.hs
index aac2f6d..e5cbd5c 100644
--- a/sky-node/lib/Config.hs
+++ b/sky-node/lib/Config.hs
@@ -8,6 +8,7 @@ import GHC.Generics (Generic)
 
 -- | Node configuration
 data AppConfig = AppConfig
-  { port :: String
+  { port :: String,
+    hostname :: String
   }
   deriving (Show, Generic, FromJSON)
diff --git a/sky-node/lib/Types.hs b/sky-node/lib/Types.hs
new file mode 100644
index 0000000..e69de29
diff --git a/sky-node/sky-node.cabal b/sky-node/sky-node.cabal
index f1c0e4e..06f326b 100644
--- a/sky-node/sky-node.cabal
+++ b/sky-node/sky-node.cabal
@@ -73,6 +73,7 @@ library sky-lib
         mtl,
         exceptions,
         stm,
+        data-default,
 
     hs-source-dirs: lib
 

From 7faa1efbc38688740218cb17e1c8451dc01a1f97 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Sun, 17 Nov 2024 22:15:27 +0200
Subject: [PATCH 06/31] feat: accept loop

---
 sky-node/config/default.yaml |  2 +-
 sky-node/lib/App.hs          | 19 +++++++++++++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/sky-node/config/default.yaml b/sky-node/config/default.yaml
index b1df009..af89220 100644
--- a/sky-node/config/default.yaml
+++ b/sky-node/config/default.yaml
@@ -1,2 +1,2 @@
 port: "3456"
-address: "localhost"
+hostname: "localhost"
diff --git a/sky-node/lib/App.hs b/sky-node/lib/App.hs
index 8b14885..19eb1e4 100644
--- a/sky-node/lib/App.hs
+++ b/sky-node/lib/App.hs
@@ -5,9 +5,12 @@ module App (runApp) where
 
 -- import App.Error
 import Config (AppConfig (..))
+import Control.Concurrent (forkIO)
 import Control.Concurrent.STM (newTVarIO)
 import Control.Concurrent.STM.TVar (TVar)
 import Control.Exception (bracket)
+import Control.Exception.Base (finally)
+import Control.Monad (when)
 import Data.Default
 import qualified Network.Socket as S
 import System.Log.FastLogger (LoggerSet, pushLogStrLn, toLogStr)
@@ -20,7 +23,7 @@ newtype Peer = Peer
   deriving (Show)
 
 -- | Signalizes if the node needs to shutdown or continue running.
-data Shutdown = Continue | Shutdown deriving (Show)
+data Shutdown = Continue | Shutdown deriving (Show, Eq, Bounded)
 
 instance Default Shutdown where
   def = Continue
@@ -69,4 +72,16 @@ runServer config state logger = do
 
 -- | Loop handling new connections.
 acceptLoop :: S.Socket -> AppState -> LoggerSet -> IO ()
-acceptLoop = undefined
+acceptLoop sock AppState {..} logger = do
+  (conn, conn_addr) <- S.accept sock
+  pushLogStrLn logger . toLogStr $ "Accepted new connection from " <> show conn_addr
+  _ <-
+    forkIO $
+      handlePeer conn AppState {..} logger `finally` do
+        S.close conn
+
+  -- loop until told to shutdown
+  when (continue == Continue) $ acceptLoop sock AppState {..} logger
+
+handlePeer :: S.Socket -> AppState -> LoggerSet -> IO ()
+handlePeer = undefined

From f2d8de033cd161474a54e042f3120acf37387ca1 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 18 Nov 2024 19:08:22 +0200
Subject: [PATCH 07/31] feat: readert pattern

---
 sky-node/app/Main.hs      |   4 +-
 sky-node/lib/App.hs       | 107 +++++++++++++++-----------------------
 sky-node/lib/App/Env.hs   |  27 ++++++++++
 sky-node/lib/App/Error.hs |   2 +-
 sky-node/lib/Peer.hs      |   8 +++
 sky-node/sky-node.cabal   |  12 +++--
 6 files changed, 89 insertions(+), 71 deletions(-)
 create mode 100644 sky-node/lib/App/Env.hs
 create mode 100644 sky-node/lib/Peer.hs

diff --git a/sky-node/app/Main.hs b/sky-node/app/Main.hs
index d5e2e31..ddf5376 100644
--- a/sky-node/app/Main.hs
+++ b/sky-node/app/Main.hs
@@ -2,7 +2,7 @@
 
 module Main where
 
-import App (runApp)
+import App (initApp)
 import Data.Yaml.Config (loadYamlSettingsArgs, useEnv)
 import System.Log.FastLogger (defaultBufSize, newStdoutLoggerSet)
 
@@ -10,4 +10,4 @@ main :: IO ()
 main = do
   logger <- newStdoutLoggerSet defaultBufSize
   config <- loadYamlSettingsArgs ["config/default.yaml"] useEnv
-  runApp config logger
+  initApp config logger
diff --git a/sky-node/lib/App.hs b/sky-node/lib/App.hs
index 19eb1e4..f35cd1b 100644
--- a/sky-node/lib/App.hs
+++ b/sky-node/lib/App.hs
@@ -1,87 +1,64 @@
 {-# LANGUAGE OverloadedRecordDot #-}
-{-# LANGUAGE RecordWildCards #-}
 
-module App (runApp) where
+module App (runApp, initApp) where
 
--- import App.Error
+import App.Env
 import Config (AppConfig (..))
-import Control.Concurrent (forkIO)
-import Control.Concurrent.STM (newTVarIO)
-import Control.Concurrent.STM.TVar (TVar)
-import Control.Exception (bracket)
-import Control.Exception.Base (finally)
-import Control.Monad (when)
-import Data.Default
+import Control.Monad.Reader (ReaderT (runReaderT), asks, liftIO)
+import Data.Functor (void)
 import qualified Network.Socket as S
 import System.Log.FastLogger (LoggerSet, pushLogStrLn, toLogStr)
+import UnliftIO.Concurrent (forkIO)
+import UnliftIO.Exception (bracket, finally)
 
--- | Aux data structure to keep info about 'Peer'.
-newtype Peer = Peer
-  { -- | Id of the 'Peer'
-    id :: String
-  }
-  deriving (Show)
+type AppM = ReaderT AppEnv IO
 
--- | Signalizes if the node needs to shutdown or continue running.
-data Shutdown = Continue | Shutdown deriving (Show, Eq, Bounded)
-
-instance Default Shutdown where
-  def = Continue
-
--- | State of the node.
-data AppState = AppState
-  { -- | Messages.
-    messages :: TVar [Int],
-    -- | List of peers that are connected to this node.
-    peers :: [Peer],
-    -- | Marker for graceful shutdown.
-    continue :: Shutdown
-  }
-
--- | Initializes default node state.
-initAppState :: IO AppState
-initAppState = do
-  let peers = def
-      continue = def
-  messages <- newTVarIO def
-  pure $ AppState {..}
+runApp :: AppEnv -> AppM a -> IO a
+runApp = flip runReaderT
 
 -- | Runs the application with the default 'AppState'.
-runApp :: AppConfig -> LoggerSet -> IO ()
-runApp config logger = do
-  state <- initAppState
-  pushLogStrLn logger $ toLogStr "Starting Sky Node..."
-  runServer config state logger
+initApp :: AppConfig -> LoggerSet -> IO ()
+initApp config logger = do
+  env <- initAppEnv config logger
+  runApp env runServer
+
+logMsg :: String -> AppM ()
+logMsg msg = do
+  logger <- asks envLogger
+  liftIO . pushLogStrLn logger $ toLogStr msg
 
 -- | Runs the server. Binds socket to the address and accepts incoming connection.
-runServer :: AppConfig -> AppState -> LoggerSet -> IO ()
-runServer config state logger = do
+runServer :: AppM ()
+runServer = do
+  logMsg "Starting Sky Node..."
+  config <- asks envConfig
   addr <- resolve config.hostname config.port
-  bracket (S.openSocket addr) S.close $ \sock -> do
-    S.setSocketOption sock S.ReuseAddr 1 -- easier for debugging
-    S.bind sock (S.addrAddress addr)
-    S.listen sock 10
-    pushLogStrLn logger . toLogStr $ "Node listening on port " <> config.port
-    acceptLoop sock state logger
+  bracket (liftIO $ S.openSocket addr) (liftIO . S.close) $ \sock -> do
+    liftIO $ do
+      S.setSocketOption sock S.ReuseAddr 1 -- easier for debugging
+      S.bind sock (S.addrAddress addr)
+      S.listen sock 10
+    logMsg $ "Node listening on port " <> config.port
+    acceptLoop sock
   where
-    resolve :: String -> String -> IO S.AddrInfo
-    resolve host port = do
+    resolve :: String -> String -> AppM S.AddrInfo
+    resolve host p = do
       let hints = S.defaultHints {S.addrSocketType = S.Stream}
-      addr : _ <- S.getAddrInfo (Just hints) (Just $ host <> ":" <> port) Nothing
+      addr : _ <- liftIO $ S.getAddrInfo (Just hints) (Just $ host <> ":" <> p) Nothing
       pure addr
 
 -- | Loop handling new connections.
-acceptLoop :: S.Socket -> AppState -> LoggerSet -> IO ()
-acceptLoop sock AppState {..} logger = do
-  (conn, conn_addr) <- S.accept sock
-  pushLogStrLn logger . toLogStr $ "Accepted new connection from " <> show conn_addr
-  _ <-
+acceptLoop :: S.Socket -> AppM ()
+acceptLoop sock = do
+  (conn, conn_addr) <- liftIO $ S.accept sock
+  logMsg $ "Accepted new connection from " <> show conn_addr
+  void $
     forkIO $
-      handlePeer conn AppState {..} logger `finally` do
-        S.close conn
+      handlePeer conn `finally` do
+        liftIO $ S.close conn
 
-  -- loop until told to shutdown
-  when (continue == Continue) $ acceptLoop sock AppState {..} logger
+-- loop until told to shutdown
+-- when (continue == Continue) $ acceptLoop sock
 
-handlePeer :: S.Socket -> AppState -> LoggerSet -> IO ()
+handlePeer :: S.Socket -> AppM ()
 handlePeer = undefined
diff --git a/sky-node/lib/App/Env.hs b/sky-node/lib/App/Env.hs
new file mode 100644
index 0000000..230a7c6
--- /dev/null
+++ b/sky-node/lib/App/Env.hs
@@ -0,0 +1,27 @@
+{-# LANGUAGE RecordWildCards #-}
+
+module App.Env
+  ( AppEnv (..),
+    initAppEnv,
+  )
+where
+
+import Config
+import Control.Concurrent.STM (TVar, newTVarIO)
+import Data.Default
+import Peer
+import System.Log.FastLogger (LoggerSet)
+
+data AppEnv = AppEnv
+  { envConfig :: AppConfig,
+    envMessages :: TVar [Int],
+    envPeers :: TVar [Peer],
+    envLogger :: LoggerSet
+  }
+
+-- | Initializes default node state.
+initAppEnv :: AppConfig -> LoggerSet -> IO AppEnv
+initAppEnv envConfig envLogger = do
+  envPeers <- newTVarIO def
+  envMessages <- newTVarIO def
+  pure $ AppEnv {..}
diff --git a/sky-node/lib/App/Error.hs b/sky-node/lib/App/Error.hs
index cfa76b9..ff860c8 100644
--- a/sky-node/lib/App/Error.hs
+++ b/sky-node/lib/App/Error.hs
@@ -2,7 +2,7 @@ module App.Error where
 
 -- | Top-level error type
 data AppError
-  = -- | Constructor for error propagation from the handler level. Used for errors that must terminate the execution of the node.
+  = -- | Constructor for error propagation from the handler level.
     HandlerError String
   | -- | Constructor for errors in the node itself.
     AppError String
diff --git a/sky-node/lib/Peer.hs b/sky-node/lib/Peer.hs
new file mode 100644
index 0000000..de9bab8
--- /dev/null
+++ b/sky-node/lib/Peer.hs
@@ -0,0 +1,8 @@
+module Peer where
+
+-- | Aux data structure to keep info about 'Peer'.
+newtype Peer = Peer
+  { -- | Id of the 'Peer'
+    id :: String
+  }
+  deriving (Show)
diff --git a/sky-node/sky-node.cabal b/sky-node/sky-node.cabal
index 06f326b..23e9023 100644
--- a/sky-node/sky-node.cabal
+++ b/sky-node/sky-node.cabal
@@ -54,22 +54,28 @@ extra-doc-files:    CHANGELOG.md
 common warnings
     ghc-options: -Wall
 
-common shared-libs
+common base
+    default-language: GHC2021
     build-depends:
         base ^>=4.18.2.1,
         yaml,
         fast-logger
 
 library sky-lib
-    import: shared-libs, warnings
+    import: base, warnings
 
     exposed-modules:
         App
         App.Error
+        App.Env
         Config
+        Peer
 
     build-depends:
         network,
+        -- effectful,
+        -- resourcet,
+        unliftio,
         mtl,
         exceptions,
         stm,
@@ -81,7 +87,7 @@ library sky-lib
 
 executable sky-node
     -- Import common warning flags.
-    import: shared-libs, warnings
+    import: base, warnings
 
     -- .hs or .lhs file containing the Main module.
     main-is:          Main.hs

From ac6813cca1d0cfbd781097790927b462bcaeaec4 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Sat, 30 Nov 2024 14:56:19 +0200
Subject: [PATCH 08/31] feat: integrating effectful[wip]

---
 sky-node/lib/App.hs           | 33 ++++++++-----------
 sky-node/lib/App/Env.hs       | 12 +++----
 sky-node/lib/Peer.hs          |  5 ++-
 sky-node/lib/Peer/Handlers.hs | 60 +++++++++++++++++++++++++++++++++++
 sky-node/lib/Types.hs         | 24 ++++++++++++++
 sky-node/sky-node.cabal       | 17 ++++++++--
 6 files changed, 121 insertions(+), 30 deletions(-)
 create mode 100644 sky-node/lib/Peer/Handlers.hs

diff --git a/sky-node/lib/App.hs b/sky-node/lib/App.hs
index f35cd1b..96e1ad1 100644
--- a/sky-node/lib/App.hs
+++ b/sky-node/lib/App.hs
@@ -1,31 +1,26 @@
+{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE OverloadedRecordDot #-}
 
-module App (runApp, initApp) where
+module App (runApp) where
 
 import App.Env
 import Config (AppConfig (..))
-import Control.Monad.Reader (ReaderT (runReaderT), asks, liftIO)
 import Data.Functor (void)
+import Effectful
+import Effectful.Concurrent
+import Effectful.Log
+import Effectful.Reader.Static
 import qualified Network.Socket as S
-import System.Log.FastLogger (LoggerSet, pushLogStrLn, toLogStr)
 import UnliftIO.Concurrent (forkIO)
 import UnliftIO.Exception (bracket, finally)
 
-type AppM = ReaderT AppEnv IO
+type AppEffects = '[Reader AppEnv, Log, Concurrent, IOE]
 
-runApp :: AppEnv -> AppM a -> IO a
-runApp = flip runReaderT
-
--- | Runs the application with the default 'AppState'.
-initApp :: AppConfig -> LoggerSet -> IO ()
-initApp config logger = do
-  env <- initAppEnv config logger
-  runApp env runServer
-
-logMsg :: String -> AppM ()
-logMsg msg = do
-  logger <- asks envLogger
-  liftIO . pushLogStrLn logger $ toLogStr msg
+runApp :: Eff AppEffects a -> IO a
+runApp = do
+  env <- initAppEnv
+  runEff $ do
+    runLogStdout $ runConcurrent $ runReader env runServer
 
 -- | Runs the server. Binds socket to the address and accepts incoming connection.
 runServer :: AppM ()
@@ -55,10 +50,8 @@ acceptLoop sock = do
   void $
     forkIO $
       handlePeer conn `finally` do
+        logMsg $ "Closing connection to " <> show conn_addr
         liftIO $ S.close conn
 
--- loop until told to shutdown
--- when (continue == Continue) $ acceptLoop sock
-
 handlePeer :: S.Socket -> AppM ()
 handlePeer = undefined
diff --git a/sky-node/lib/App/Env.hs b/sky-node/lib/App/Env.hs
index 230a7c6..3959712 100644
--- a/sky-node/lib/App/Env.hs
+++ b/sky-node/lib/App/Env.hs
@@ -9,19 +9,19 @@ where
 import Config
 import Control.Concurrent.STM (TVar, newTVarIO)
 import Data.Default
+import Data.IntMap.Strict (IntMap)
 import Peer
-import System.Log.FastLogger (LoggerSet)
+import Types
 
 data AppEnv = AppEnv
   { envConfig :: AppConfig,
-    envMessages :: TVar [Int],
-    envPeers :: TVar [Peer],
-    envLogger :: LoggerSet
+    envMessages :: TVar (IntMap Topic),
+    envPeers :: TVar [Peer]
   }
 
 -- | Initializes default node state.
-initAppEnv :: AppConfig -> LoggerSet -> IO AppEnv
-initAppEnv envConfig envLogger = do
+initAppEnv :: AppConfig -> IO AppEnv
+initAppEnv envConfig = do
   envPeers <- newTVarIO def
   envMessages <- newTVarIO def
   pure $ AppEnv {..}
diff --git a/sky-node/lib/Peer.hs b/sky-node/lib/Peer.hs
index de9bab8..1e13491 100644
--- a/sky-node/lib/Peer.hs
+++ b/sky-node/lib/Peer.hs
@@ -1,4 +1,7 @@
-module Peer where
+module Peer
+  ( Peer (..),
+  )
+where
 
 -- | Aux data structure to keep info about 'Peer'.
 newtype Peer = Peer
diff --git a/sky-node/lib/Peer/Handlers.hs b/sky-node/lib/Peer/Handlers.hs
new file mode 100644
index 0000000..20b888b
--- /dev/null
+++ b/sky-node/lib/Peer/Handlers.hs
@@ -0,0 +1,60 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+
+module Peer.Handlers where
+
+import App.Env
+import Control.Concurrent.STM (atomically, readTVar, writeTVar)
+import qualified Data.IntMap.Strict as IntMap
+import Data.Word (Word64)
+import Effectful
+import Effectful.Dispatch.Dynamic
+import Effectful.Error.Dynamic
+import Effectful.Reader.Dynamic (Reader, asks)
+import Effectful.TH (makeEffect)
+import Types
+import UnliftIO (catchIO, throwString)
+
+data PeerHandler :: Effect where
+  PublishBlock :: TopicId -> BlockData -> PeerHandler m (Certificate Block)
+  GetTopics :: Maybe TopicId -> PeerHandler m [TopicId]
+  DescribeTopic :: TopicId -> PeerHandler m TopicMetaData
+  PollTopic :: TopicId -> PeerHandler m (Word64, Certificate Block)
+  GetTopicBlockCertificate :: TopicId -> Word64 -> PeerHandler m (Certificate Block)
+  ReadTopic :: TopicId -> Word64 -> PeerHandler m [BlockData]
+
+-- TODO replace with makeEffect_ and add docs for type sigs, or add docs to the GADT
+makeEffect ''PeerHandler
+
+runPeerHandlerIO :: (IOE :> es, Error String :> es, Reader AppEnv :> es) => Eff (PeerHandler : es) a -> Eff es a
+runPeerHandlerIO = interpret $ \_ -> \case
+  PublishBlock tId bData -> do
+    topics <- asks envMessages
+    adapt $ atomically $ do
+      tpcs <- readTVar topics
+      let tpcs' = IntMap.adjust (\tpc -> tpc {topicMessages = IntMap.insert (IntMap.size (topicMessages tpc)) bData (topicMessages tpc)}) (undefined {- this must be TopicId -}) tpcs
+      writeTVar topics tpcs'
+    pure $ makeBlockCertificate bData
+  GetTopics mtId -> adapt $ undefined
+  DescribeTopic tId -> do
+    topics <- asks envMessages
+    adapt $ do
+      maybeMeta <- atomically $ do
+        tpcs <- readTVar topics
+        let mTpc = topicMeta <$> (IntMap.!?) tpcs (undefined {- this must be TopicId -})
+        pure mTpc
+      case maybeMeta of
+        Just meta -> pure meta
+        Nothing -> throwString "No such topic"
+  PollTopic tId -> adapt $ undefined
+  GetTopicBlockCertificate tId height -> adapt $ undefined
+  ReadTopic tId height -> adapt $ undefined
+  where
+    adapt m = liftIO m `catchIO` \e -> throwString $ show e
diff --git a/sky-node/lib/Types.hs b/sky-node/lib/Types.hs
index e69de29..aff1a88 100644
--- a/sky-node/lib/Types.hs
+++ b/sky-node/lib/Types.hs
@@ -0,0 +1,24 @@
+module Types where
+
+import qualified Data.ByteString as BS
+import Data.IntMap.Strict (IntMap)
+
+newtype TopicId = TopicId {topicId :: BS.ByteString} deriving (Show, Eq)
+
+newtype TopicMetaData = TopicMetaData {topicMetaId :: TopicId} deriving (Show, Eq)
+
+data Topic = Topic
+  { topicMeta :: TopicMetaData,
+    topicMessages :: IntMap BlockData
+  }
+  deriving (Show, Eq)
+
+data Block
+
+newtype BlockData = BlockData {blockData :: BS.ByteString} deriving (Show, Eq)
+
+newtype Certificate a = Certificate {cert :: BS.ByteString} deriving (Show, Eq)
+
+-- TODO for now it's a stub
+makeBlockCertificate :: BlockData -> Certificate Block
+makeBlockCertificate = undefined
diff --git a/sky-node/sky-node.cabal b/sky-node/sky-node.cabal
index 23e9023..1f18c17 100644
--- a/sky-node/sky-node.cabal
+++ b/sky-node/sky-node.cabal
@@ -59,7 +59,7 @@ common base
     build-depends:
         base ^>=4.18.2.1,
         yaml,
-        fast-logger
+        log-base,
 
 library sky-lib
     import: base, warnings
@@ -70,17 +70,28 @@ library sky-lib
         App.Env
         Config
         Peer
+        Peer.Handlers
+        Types
 
     build-depends:
         network,
-        -- effectful,
+        effectful,
+        effectful-core,
+        log-effectful,
+        effectful-plugin,
+        effectful-th,
+        async,
+        bytestring,
+        containers,
         -- resourcet,
+        -- mtl,
         unliftio,
-        mtl,
         exceptions,
         stm,
         data-default,
 
+    ghc-options: -fplugin=Effectful.Plugin
+
     hs-source-dirs: lib
 
     default-language: Haskell2010

From dc72bf6f555587c15e733aea0851424b568b3948 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Sat, 30 Nov 2024 15:49:43 +0200
Subject: [PATCH 09/31] feat: integrating effectful

---
 sky-node/app/Main.hs          |  6 +--
 sky-node/lib/App.hs           | 77 +++++++++++++++++------------------
 sky-node/lib/Config.hs        |  5 ++-
 sky-node/lib/Peer/Handlers.hs |  9 ++--
 sky-node/lib/Types.hs         |  2 +-
 sky-node/sky-node.cabal       |  1 +
 6 files changed, 51 insertions(+), 49 deletions(-)

diff --git a/sky-node/app/Main.hs b/sky-node/app/Main.hs
index ddf5376..c1cc5c0 100644
--- a/sky-node/app/Main.hs
+++ b/sky-node/app/Main.hs
@@ -4,10 +4,10 @@ module Main where
 
 import App (initApp)
 import Data.Yaml.Config (loadYamlSettingsArgs, useEnv)
-import System.Log.FastLogger (defaultBufSize, newStdoutLoggerSet)
+import Log.Backend.StandardOutput
 
 main :: IO ()
 main = do
-  logger <- newStdoutLoggerSet defaultBufSize
   config <- loadYamlSettingsArgs ["config/default.yaml"] useEnv
-  initApp config logger
+  withStdOutLogger $ \stdoutLogger -> do
+    initApp config stdoutLogger
diff --git a/sky-node/lib/App.hs b/sky-node/lib/App.hs
index 96e1ad1..7b8c8b1 100644
--- a/sky-node/lib/App.hs
+++ b/sky-node/lib/App.hs
@@ -1,57 +1,56 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE OverloadedStrings #-}
 
-module App (runApp) where
+module App (initApp) where
 
 import App.Env
 import Config (AppConfig (..))
-import Data.Functor (void)
+import Control.Monad (forever)
+import Data.Text (pack, unpack)
 import Effectful
 import Effectful.Concurrent
+import Effectful.Concurrent.Async
 import Effectful.Log
 import Effectful.Reader.Static
 import qualified Network.Socket as S
-import UnliftIO.Concurrent (forkIO)
-import UnliftIO.Exception (bracket, finally)
+import UnliftIO.Exception (bracket)
 
-type AppEffects = '[Reader AppEnv, Log, Concurrent, IOE]
+type AppEffects = '[Reader AppEnv, Concurrent, Log, IOE]
 
-runApp :: Eff AppEffects a -> IO a
-runApp = do
-  env <- initAppEnv
+initApp :: AppConfig -> Logger -> IO ()
+initApp config logger = do
+  env <- initAppEnv config
   runEff $ do
-    runLogStdout $ runConcurrent $ runReader env runServer
+    runLog "main" logger defaultLogLevel $
+      runConcurrent $
+        runReader env $ do
+          withSocketServer handlePeer
 
--- | Runs the server. Binds socket to the address and accepts incoming connection.
-runServer :: AppM ()
-runServer = do
-  logMsg "Starting Sky Node..."
+withSocketServer :: (S.Socket -> Eff AppEffects ()) -> Eff AppEffects ()
+withSocketServer handler = do
+  logInfo_ "Starting Sky Node..."
   config <- asks envConfig
-  addr <- resolve config.hostname config.port
-  bracket (liftIO $ S.openSocket addr) (liftIO . S.close) $ \sock -> do
-    liftIO $ do
-      S.setSocketOption sock S.ReuseAddr 1 -- easier for debugging
-      S.bind sock (S.addrAddress addr)
-      S.listen sock 10
-    logMsg $ "Node listening on port " <> config.port
-    acceptLoop sock
-  where
-    resolve :: String -> String -> AppM S.AddrInfo
-    resolve host p = do
-      let hints = S.defaultHints {S.addrSocketType = S.Stream}
-      addr : _ <- liftIO $ S.getAddrInfo (Just hints) (Just $ host <> ":" <> p) Nothing
-      pure addr
+  bracket setupServerSocket (liftIO . S.close) $ \sock -> do
+    logInfo_ $ "Node listening on " <> config.hostname <> ":" <> config.port
+    forever $ do
+      (conn, addr) <- liftIO $ S.accept sock
+      logInfo_ $ "Connection accepted from " <> pack (show addr)
+      -- using withAsync ensures proper thread handling in face of execption
+      withAsync (handler conn) $ \_ -> pure ()
 
--- | Loop handling new connections.
-acceptLoop :: S.Socket -> AppM ()
-acceptLoop sock = do
-  (conn, conn_addr) <- liftIO $ S.accept sock
-  logMsg $ "Accepted new connection from " <> show conn_addr
-  void $
-    forkIO $
-      handlePeer conn `finally` do
-        logMsg $ "Closing connection to " <> show conn_addr
-        liftIO $ S.close conn
+setupServerSocket :: Eff AppEffects S.Socket
+setupServerSocket = do
+  h <- asks $ hostname . envConfig
+  p <- asks $ port . envConfig
+  addrs <- liftIO $ S.getAddrInfo (Just S.defaultHints {S.addrSocketType = S.Stream}) (Just $ unpack h) (Just $ unpack p)
+  let addr = head addrs
+  sock <- liftIO $ S.openSocket addr
+  liftIO $ do
+    S.setSocketOption sock S.ReuseAddr 1 -- easier for debugging
+    S.bind sock (S.addrAddress addr)
+    S.listen sock 10
+  pure sock
 
-handlePeer :: S.Socket -> AppM ()
-handlePeer = undefined
+handlePeer :: S.Socket -> Eff AppEffects ()
+handlePeer sock = undefined
diff --git a/sky-node/lib/Config.hs b/sky-node/lib/Config.hs
index e5cbd5c..979b157 100644
--- a/sky-node/lib/Config.hs
+++ b/sky-node/lib/Config.hs
@@ -3,12 +3,13 @@
 
 module Config where
 
+import Data.Text (Text)
 import Data.Yaml (FromJSON)
 import GHC.Generics (Generic)
 
 -- | Node configuration
 data AppConfig = AppConfig
-  { port :: String,
-    hostname :: String
+  { port :: Text,
+    hostname :: Text
   }
   deriving (Show, Generic, FromJSON)
diff --git a/sky-node/lib/Peer/Handlers.hs b/sky-node/lib/Peer/Handlers.hs
index 20b888b..ab157fc 100644
--- a/sky-node/lib/Peer/Handlers.hs
+++ b/sky-node/lib/Peer/Handlers.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE TypeApplications #-}
@@ -35,20 +36,20 @@ makeEffect ''PeerHandler
 
 runPeerHandlerIO :: (IOE :> es, Error String :> es, Reader AppEnv :> es) => Eff (PeerHandler : es) a -> Eff es a
 runPeerHandlerIO = interpret $ \_ -> \case
-  PublishBlock tId bData -> do
+  PublishBlock TopicId {..} bData -> do
     topics <- asks envMessages
     adapt $ atomically $ do
       tpcs <- readTVar topics
-      let tpcs' = IntMap.adjust (\tpc -> tpc {topicMessages = IntMap.insert (IntMap.size (topicMessages tpc)) bData (topicMessages tpc)}) (undefined {- this must be TopicId -}) tpcs
+      let tpcs' = IntMap.adjust (\tpc -> tpc {topicMessages = IntMap.insert (IntMap.size (topicMessages tpc)) bData (topicMessages tpc)}) topicId tpcs
       writeTVar topics tpcs'
     pure $ makeBlockCertificate bData
   GetTopics mtId -> adapt $ undefined
-  DescribeTopic tId -> do
+  DescribeTopic TopicId {..} -> do
     topics <- asks envMessages
     adapt $ do
       maybeMeta <- atomically $ do
         tpcs <- readTVar topics
-        let mTpc = topicMeta <$> (IntMap.!?) tpcs (undefined {- this must be TopicId -})
+        let mTpc = topicMeta <$> (IntMap.!?) tpcs topicId
         pure mTpc
       case maybeMeta of
         Just meta -> pure meta
diff --git a/sky-node/lib/Types.hs b/sky-node/lib/Types.hs
index aff1a88..21fb8e4 100644
--- a/sky-node/lib/Types.hs
+++ b/sky-node/lib/Types.hs
@@ -3,7 +3,7 @@ module Types where
 import qualified Data.ByteString as BS
 import Data.IntMap.Strict (IntMap)
 
-newtype TopicId = TopicId {topicId :: BS.ByteString} deriving (Show, Eq)
+newtype TopicId = TopicId {topicId :: Int} deriving (Show, Eq)
 
 newtype TopicMetaData = TopicMetaData {topicMetaId :: TopicId} deriving (Show, Eq)
 
diff --git a/sky-node/sky-node.cabal b/sky-node/sky-node.cabal
index 1f18c17..cc36832 100644
--- a/sky-node/sky-node.cabal
+++ b/sky-node/sky-node.cabal
@@ -81,6 +81,7 @@ library sky-lib
         effectful-plugin,
         effectful-th,
         async,
+        text,
         bytestring,
         containers,
         -- resourcet,

From 9d44618d7618ee1fd0d89f3bdace477f67ac2729 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Sat, 30 Nov 2024 17:14:15 +0200
Subject: [PATCH 10/31] feat: package for merkle patricia trie in haskell

---
 merkle-patricia-trie/.envrc                   |  2 +
 merkle-patricia-trie/CHANGELOG.md             |  5 ++
 merkle-patricia-trie/flake.lock               | 61 +++++++++++++++
 merkle-patricia-trie/flake.nix                | 30 ++++++++
 merkle-patricia-trie/lib/Data/MerkleTrie.hs   |  1 +
 .../merkle-patricia-trie.cabal                | 75 +++++++++++++++++++
 6 files changed, 174 insertions(+)
 create mode 100644 merkle-patricia-trie/.envrc
 create mode 100644 merkle-patricia-trie/CHANGELOG.md
 create mode 100644 merkle-patricia-trie/flake.lock
 create mode 100644 merkle-patricia-trie/flake.nix
 create mode 100644 merkle-patricia-trie/lib/Data/MerkleTrie.hs
 create mode 100644 merkle-patricia-trie/merkle-patricia-trie.cabal

diff --git a/merkle-patricia-trie/.envrc b/merkle-patricia-trie/.envrc
new file mode 100644
index 0000000..a971e5d
--- /dev/null
+++ b/merkle-patricia-trie/.envrc
@@ -0,0 +1,2 @@
+watch_file *.cabal
+use flake
diff --git a/merkle-patricia-trie/CHANGELOG.md b/merkle-patricia-trie/CHANGELOG.md
new file mode 100644
index 0000000..246ac3f
--- /dev/null
+++ b/merkle-patricia-trie/CHANGELOG.md
@@ -0,0 +1,5 @@
+# Revision history for merkle-patricia-trie
+
+## 0.1.0.0 -- YYYY-mm-dd
+
+* First version. Released on an unsuspecting world.
diff --git a/merkle-patricia-trie/flake.lock b/merkle-patricia-trie/flake.lock
new file mode 100644
index 0000000..0b3070f
--- /dev/null
+++ b/merkle-patricia-trie/flake.lock
@@ -0,0 +1,61 @@
+{
+  "nodes": {
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1732758367,
+        "narHash": "sha256-RzaI1RO0UXqLjydtz3GAXSTzHkpb/lLD1JD8a0W4Wpo=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "fa42b5a5f401aab8a32bd33c9a4de0738180dc59",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}
diff --git a/merkle-patricia-trie/flake.nix b/merkle-patricia-trie/flake.nix
new file mode 100644
index 0000000..d76f439
--- /dev/null
+++ b/merkle-patricia-trie/flake.nix
@@ -0,0 +1,30 @@
+{
+  inputs = {
+    flake-utils.url = "github:numtide/flake-utils";
+    nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
+  };
+  outputs = { self, nixpkgs, flake-utils, ... }: flake-utils.lib.eachDefaultSystem (system:
+    let
+      pkgs = import nixpkgs { inherit system; };
+      overlay = final: prev: {
+        merkle-patricia-trie = final.callCabal2nix "merkle-patricia-trie" ./. { };
+      };
+      haskPkgs = pkgs.haskellPackages.extend overlay;
+    in
+    {
+      devShells.default = haskPkgs.shellFor {
+        packages = p: [
+          p.merkle-patricia-trie
+        ];
+        nativeBuildInputs = with haskPkgs; [
+          cabal-install
+          haskell-language-server
+          hlint
+        ];
+        shellHook = ''
+          echo "Welcome"
+        '';
+      };
+    }
+  );
+}
diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
new file mode 100644
index 0000000..5a79ff0
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -0,0 +1 @@
+module Data.MerkleTrie where
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
new file mode 100644
index 0000000..e018132
--- /dev/null
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -0,0 +1,75 @@
+cabal-version:      3.0
+-- The cabal-version field refers to the version of the .cabal specification,
+-- and can be different from the cabal-install (the tool) version and the
+-- Cabal (the library) version you are using. As such, the Cabal (the library)
+-- version used must be equal or greater than the version stated in this field.
+-- Starting from the specification version 2.2, the cabal-version field must be
+-- the first thing in the cabal file.
+
+-- Initial package description 'merkle-patricia-trie' generated by
+-- 'cabal init'. For further documentation, see:
+--   http://haskell.org/cabal/users-guide/
+--
+-- The name of the package.
+name:               merkle-patricia-trie
+
+-- The package version.
+-- See the Haskell package versioning policy (PVP) for standards
+-- guiding when and how versions should be incremented.
+-- https://pvp.haskell.org
+-- PVP summary:     +-+------- breaking API changes
+--                  | | +----- non-breaking API additions
+--                  | | | +--- code changes with no API change
+version:            0.1.0.0
+
+-- A short (one-line) description of the package.
+synopsis:
+    Implementation of Merkle Patricia Trie data structure in Haskell
+
+-- A longer description of the package.
+-- description:
+
+-- The license under which the package is released.
+license:            NONE
+
+-- The package author(s).
+author:             Yaroslav Kozhevnikov
+
+-- An email address to which users can send suggestions, bug reports, and patches.
+maintainer:         yroslav541@gmail.com
+
+-- A copyright notice.
+-- copyright:
+category:           Data
+build-type:         Simple
+
+-- Extra doc files to be distributed with the package, such as a CHANGELOG or a README.
+extra-doc-files:    CHANGELOG.md
+
+-- Extra source files to be distributed with the package, such as examples, or a tutorial module.
+-- extra-source-files:
+
+common warnings
+    ghc-options: -Wall
+
+library
+    -- Import common warning flags.
+    import:           warnings
+
+    -- Modules exported by the library.
+    exposed-modules:  Data.MerkleTrie
+
+    -- Modules included in this library but not exported.
+    -- other-modules:
+
+    -- LANGUAGE extensions used by modules in this package.
+    -- other-extensions:
+
+    -- Other library packages from which modules are imported.
+    build-depends:    base ^>=4.18.2.1
+
+    -- Directories containing source files.
+    hs-source-dirs:   lib
+
+    -- Base language which the package is written in.
+    default-language: Haskell2010

From 060d9722155e567beef6a45bae7f59c5aca802e9 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Sat, 30 Nov 2024 18:13:57 +0200
Subject: [PATCH 11/31] feat: simple naive implementation of Merkle Patricia
 Trie.\nTo be optimized

---
 merkle-patricia-trie/lib/Data/MerkleTrie.hs   | 102 ++++++++++++++++++
 .../merkle-patricia-trie.cabal                |   7 +-
 2 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index 5a79ff0..b1deaf9 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -1 +1,103 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+
 module Data.MerkleTrie where
+
+import Control.Applicative ((<|>))
+import Crypto.Hash (Digest, SHA256, hash)
+import Data.ByteString.Char8 (pack)
+import Data.List (isPrefixOf)
+import GHC.Generics (Generic)
+
+-- | Patricia Trie node
+data PatriciaTrie k v = Node
+  { prefix :: k,
+    value :: Maybe v,
+    children :: [PatriciaTrie k v]
+  }
+  deriving (Show, Generic)
+
+-- | Merkle Patricia Trie
+data MerklePatriciaTrie k v = MerklePatriciaTrie
+  { rootHash :: Digest SHA256,
+    rootNode :: PatriciaTrie k v
+  }
+  deriving (Show, Generic)
+
+-- | Compute the hash of a Patricia Trie node
+computeNodeHash :: (Show k, Show v) => PatriciaTrie k v -> Digest SHA256
+computeNodeHash Node {..} =
+  hash . pack $ show (prefix, value, map computeNodeHash children)
+
+-- | Update the root hash of the Merkle Trie
+updateRootHash :: (Show k, Show v) => PatriciaTrie k v -> Digest SHA256
+updateRootHash = computeNodeHash
+
+-- | Create an empty Merkle Patricia Trie
+emptyTrie :: forall k v. (Show k, Show v) => MerklePatriciaTrie [k] v
+emptyTrie =
+  MerklePatriciaTrie
+    { rootHash = computeNodeHash @[k] @v root,
+      rootNode = root
+    }
+  where
+    root = Node {prefix = [], value = Nothing, children = []}
+
+-- | Create a singleton Patricia Trie
+singletonTrie :: forall k v. (Show k, Show v, Eq k) => [k] -> v -> MerklePatriciaTrie [k] v
+singletonTrie k v =
+  MerklePatriciaTrie
+    { rootHash = computeNodeHash @[k] @v root,
+      rootNode = root
+    }
+  where
+    root = Node {prefix = k, value = Just v, children = []}
+
+-- | Insert a key-value pair into a Patricia Trie node
+insertNode :: (Eq k, Ord k) => [k] -> v -> PatriciaTrie [k] v -> PatriciaTrie [k] v
+insertNode [] v node = node {value = Just v}
+insertNode key v Node {..}
+  | prefix `isPrefixOf` key =
+      let remaining = drop (length prefix) key
+       in if null remaining
+            then Node {value = Just v, ..}
+            else Node {children = insertChild remaining v children, ..}
+  | otherwise = mergeNodes key v prefix value children
+  where
+    insertChild [] val [] = [Node {prefix = [], value = Just val, children = []}]
+    insertChild r val (c : cs)
+      | commonPrefix r c.prefix /= [] =
+          let splitChild = insertNode (drop (length (commonPrefix r c.prefix)) r) val c
+           in splitChild : cs
+      | otherwise = c : insertChild r val cs
+    insertChild r val [] = [Node {prefix = r, value = Just val, children = []}]
+    mergeNodes newKey newValue oldKey oldValue oldChildren =
+      let common = commonPrefix newKey oldKey
+          splitOld = Node {prefix = drop (length common) oldKey, value = oldValue, children = oldChildren}
+          splitNew = Node {prefix = drop (length common) newKey, value = Just newValue, children = []}
+       in Node {prefix = common, value = Nothing, children = [splitOld, splitNew]}
+    commonPrefix :: (Eq a) => [a] -> [a] -> [a]
+    commonPrefix [] _ = []
+    commonPrefix _ [] = []
+    commonPrefix (x : xs) (y : ys)
+      | x == y = x : commonPrefix xs ys
+      | otherwise = []
+
+-- | Insert into the Merkle Patricia Trie
+insertTrie :: (Show k, Ord k, Eq k, Show v) => [k] -> v -> MerklePatriciaTrie [k] v -> MerklePatriciaTrie [k] v
+insertTrie k v MerklePatriciaTrie {rootNode} =
+  let newRoot = insertNode k v rootNode
+   in MerklePatriciaTrie {rootHash = updateRootHash newRoot, rootNode = newRoot}
+
+-- | Generate a Merkle proof for a key
+merkleProof :: forall k v. (Eq k, Show k, Show v) => [k] -> PatriciaTrie [k] v -> Maybe (v, Digest SHA256)
+merkleProof [] Node {value} = (\v -> (v, computeNodeHash @[k] @v Node {prefix = [], value = Just v, children = []})) <$> value
+merkleProof key Node {prefix, children}
+  | prefix `isPrefixOf` key =
+      let remaining = drop (length prefix) key
+       in foldr (\child acc -> acc <|> merkleProof remaining child) Nothing children
+  | otherwise = Nothing
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index e018132..ed41911 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -66,7 +66,12 @@ library
     -- other-extensions:
 
     -- Other library packages from which modules are imported.
-    build-depends:    base ^>=4.18.2.1
+    build-depends:
+        base ^>=4.18.2.1,
+        cryptonite ==0.30,
+        bytestring
+        -- for future more optimized implementation consider using RLP from ethereum whitepaper
+        -- relapse ==1.0.0.1
 
     -- Directories containing source files.
     hs-source-dirs:   lib

From 6984a7b019e30f0a2208ea6b9073dc49f767c1b2 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Sun, 1 Dec 2024 14:49:25 +0200
Subject: [PATCH 12/31] feat: some refactoring

---
 sky-node/config/default.yaml  |  2 +-
 sky-node/lib/App.hs           | 11 +++++------
 sky-node/lib/App/Env.hs       | 35 ++++++++++++++++++++++++++++-------
 sky-node/lib/Config.hs        |  5 +----
 sky-node/lib/Peer/Handlers.hs | 23 ++++++++++-------------
 sky-node/lib/Utils.hs         | 16 ++++++++++++++++
 sky-node/sky-node.cabal       | 32 ++++++++++++++++++++++++++++++++
 7 files changed, 93 insertions(+), 31 deletions(-)
 create mode 100644 sky-node/lib/Utils.hs

diff --git a/sky-node/config/default.yaml b/sky-node/config/default.yaml
index af89220..e8300aa 100644
--- a/sky-node/config/default.yaml
+++ b/sky-node/config/default.yaml
@@ -1,2 +1,2 @@
 port: "3456"
-hostname: "localhost"
+host: "localhost"
diff --git a/sky-node/lib/App.hs b/sky-node/lib/App.hs
index 7b8c8b1..d3a54d6 100644
--- a/sky-node/lib/App.hs
+++ b/sky-node/lib/App.hs
@@ -1,6 +1,4 @@
 {-# LANGUAGE DataKinds #-}
-{-# LANGUAGE OverloadedRecordDot #-}
-{-# LANGUAGE OverloadedStrings #-}
 
 module App (initApp) where
 
@@ -15,6 +13,7 @@ import Effectful.Log
 import Effectful.Reader.Static
 import qualified Network.Socket as S
 import UnliftIO.Exception (bracket)
+import Utils
 
 type AppEffects = '[Reader AppEnv, Concurrent, Log, IOE]
 
@@ -30,9 +29,9 @@ initApp config logger = do
 withSocketServer :: (S.Socket -> Eff AppEffects ()) -> Eff AppEffects ()
 withSocketServer handler = do
   logInfo_ "Starting Sky Node..."
-  config <- asks envConfig
+  config <- askFieldS @AppConfig
   bracket setupServerSocket (liftIO . S.close) $ \sock -> do
-    logInfo_ $ "Node listening on " <> config.hostname <> ":" <> config.port
+    logInfo_ $ "Node listening on " <> config.host <> ":" <> config.port
     forever $ do
       (conn, addr) <- liftIO $ S.accept sock
       logInfo_ $ "Connection accepted from " <> pack (show addr)
@@ -41,8 +40,8 @@ withSocketServer handler = do
 
 setupServerSocket :: Eff AppEffects S.Socket
 setupServerSocket = do
-  h <- asks $ hostname . envConfig
-  p <- asks $ port . envConfig
+  h <- (.host) <$> askFieldS @AppConfig
+  p <- (.port) <$> askFieldS @AppConfig
   addrs <- liftIO $ S.getAddrInfo (Just S.defaultHints {S.addrSocketType = S.Stream}) (Just $ unpack h) (Just $ unpack p)
   let addr = head addrs
   sock <- liftIO $ S.openSocket addr
diff --git a/sky-node/lib/App/Env.hs b/sky-node/lib/App/Env.hs
index 3959712..de81f8c 100644
--- a/sky-node/lib/App/Env.hs
+++ b/sky-node/lib/App/Env.hs
@@ -1,8 +1,12 @@
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE DataKinds #-}
+-- while looks scary, it's needed because GHC can't guarantee HasField instances. Just be careful when defining them.
+{-# LANGUAGE UndecidableInstances #-}
 
 module App.Env
   ( AppEnv (..),
     initAppEnv,
+    Topics,
+    Peers,
   )
 where
 
@@ -10,18 +14,35 @@ import Config
 import Control.Concurrent.STM (TVar, newTVarIO)
 import Data.Default
 import Data.IntMap.Strict (IntMap)
+import GHC.Records (HasField)
 import Peer
 import Types
+import Utils
+
+-- | Convenvince type alias to avoid writing full type in 'askField' invocations
+type Topics = TVar (IntMap Topic)
+
+-- | Convenvince type alias to avoid writing full type in 'askField' invocations
+type Peers = TVar [Peer]
 
 data AppEnv = AppEnv
-  { envConfig :: AppConfig,
-    envMessages :: TVar (IntMap Topic),
-    envPeers :: TVar [Peer]
+  { config :: AppConfig,
+    topics :: Topics,
+    peers :: Peers
   }
 
+instance (HasField "config" AppEnv AppConfig) => Has AppConfig AppEnv where
+  getField env = env.config
+
+instance (HasField "topics" AppEnv Topics) => Has Topics AppEnv where
+  getField env = env.topics
+
+instance (HasField "peers" AppEnv Peers) => Has Peers AppEnv where
+  getField env = env.peers
+
 -- | Initializes default node state.
 initAppEnv :: AppConfig -> IO AppEnv
-initAppEnv envConfig = do
-  envPeers <- newTVarIO def
-  envMessages <- newTVarIO def
+initAppEnv config = do
+  peers <- newTVarIO def
+  topics <- newTVarIO def
   pure $ AppEnv {..}
diff --git a/sky-node/lib/Config.hs b/sky-node/lib/Config.hs
index 979b157..6a957c5 100644
--- a/sky-node/lib/Config.hs
+++ b/sky-node/lib/Config.hs
@@ -1,6 +1,3 @@
-{-# LANGUAGE DeriveAnyClass #-}
-{-# LANGUAGE DeriveGeneric #-}
-
 module Config where
 
 import Data.Text (Text)
@@ -10,6 +7,6 @@ import GHC.Generics (Generic)
 -- | Node configuration
 data AppConfig = AppConfig
   { port :: Text,
-    hostname :: Text
+    host :: Text
   }
   deriving (Show, Generic, FromJSON)
diff --git a/sky-node/lib/Peer/Handlers.hs b/sky-node/lib/Peer/Handlers.hs
index ab157fc..530b971 100644
--- a/sky-node/lib/Peer/Handlers.hs
+++ b/sky-node/lib/Peer/Handlers.hs
@@ -1,27 +1,24 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell #-}
-{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
 
 module Peer.Handlers where
 
-import App.Env
+import App.Env (Topics)
 import Control.Concurrent.STM (atomically, readTVar, writeTVar)
 import qualified Data.IntMap.Strict as IntMap
 import Data.Word (Word64)
 import Effectful
 import Effectful.Dispatch.Dynamic
-import Effectful.Error.Dynamic
-import Effectful.Reader.Dynamic (Reader, asks)
+import Effectful.Error.Static
+import Effectful.Reader.Static (Reader)
 import Effectful.TH (makeEffect)
 import Types
 import UnliftIO (catchIO, throwString)
+import Utils
 
 data PeerHandler :: Effect where
   PublishBlock :: TopicId -> BlockData -> PeerHandler m (Certificate Block)
@@ -34,26 +31,26 @@ data PeerHandler :: Effect where
 -- TODO replace with makeEffect_ and add docs for type sigs, or add docs to the GADT
 makeEffect ''PeerHandler
 
-runPeerHandlerIO :: (IOE :> es, Error String :> es, Reader AppEnv :> es) => Eff (PeerHandler : es) a -> Eff es a
+runPeerHandlerIO :: (Has Topics env, IOE :> es, Error String :> es, Reader env :> es) => Eff (PeerHandler : es) a -> Eff es a
 runPeerHandlerIO = interpret $ \_ -> \case
   PublishBlock TopicId {..} bData -> do
-    topics <- asks envMessages
+    topics <- askFieldS @Topics
     adapt $ atomically $ do
       tpcs <- readTVar topics
-      let tpcs' = IntMap.adjust (\tpc -> tpc {topicMessages = IntMap.insert (IntMap.size (topicMessages tpc)) bData (topicMessages tpc)}) topicId tpcs
+      let tpcs' = IntMap.adjust (\tpc -> tpc {topicMessages = IntMap.insert (IntMap.size tpc.topicMessages) bData tpc.topicMessages}) topicId tpcs
       writeTVar topics tpcs'
     pure $ makeBlockCertificate bData
   GetTopics mtId -> adapt $ undefined
   DescribeTopic TopicId {..} -> do
-    topics <- asks envMessages
+    topics <- askFieldS @Topics
     adapt $ do
       maybeMeta <- atomically $ do
         tpcs <- readTVar topics
-        let mTpc = topicMeta <$> (IntMap.!?) tpcs topicId
+        let mTpc = (.topicMeta) <$> (IntMap.!?) tpcs topicId
         pure mTpc
       case maybeMeta of
         Just meta -> pure meta
-        Nothing -> throwString "No such topic"
+        Nothing -> throwString "No such topic" -- TODO replace with proper error handling
   PollTopic tId -> adapt $ undefined
   GetTopicBlockCertificate tId height -> adapt $ undefined
   ReadTopic tId height -> adapt $ undefined
diff --git a/sky-node/lib/Utils.hs b/sky-node/lib/Utils.hs
new file mode 100644
index 0000000..6f20436
--- /dev/null
+++ b/sky-node/lib/Utils.hs
@@ -0,0 +1,16 @@
+module Utils where
+
+import Effectful
+import qualified Effectful.Reader.Dynamic as D
+import qualified Effectful.Reader.Static as S
+
+-- | Utility class to abstract over (Reader rec) environment.
+-- Allows for easier refactoring, since you can constrain what functions need what from the environment.
+class Has field rec where
+  getField :: rec -> field
+
+askFieldS :: (Has field env, S.Reader env :> es) => Eff es field
+askFieldS = S.asks getField
+
+askFieldD :: (Has field env, D.Reader env :> es) => Eff es field
+askFieldD = D.asks getField
diff --git a/sky-node/sky-node.cabal b/sky-node/sky-node.cabal
index cc36832..b47a440 100644
--- a/sky-node/sky-node.cabal
+++ b/sky-node/sky-node.cabal
@@ -72,6 +72,7 @@ library sky-lib
         Peer
         Peer.Handlers
         Types
+        Utils
 
     build-depends:
         network,
@@ -91,6 +92,37 @@ library sky-lib
         stm,
         data-default,
 
+    default-extensions:
+        -- enables {..} syntax for records
+        RecordWildCards,
+        -- enables {a} syntax to use instead of {a = a}
+        NamedFieldPuns,
+        -- doesn't generate field selector functions
+        NoFieldSelectors,
+        -- allows to have multiple fields with the same name(but no same name fields in one constructor!)
+        DuplicateRecordFields,
+        -- disambiguates between the same name fields in different records
+        DisambiguateRecordFields,
+        -- enables a.b syntax for records
+        OverloadedRecordDot,
+        -- allows to use "string" literals to get different string types(i.e. Text, String, ByteString)
+        OverloadedStrings,
+        -- some relaxing of GHC constraints on typeclass declaration/resolution
+        FlexibleContexts,
+        FlexibleInstances,
+        MultiParamTypeClasses,
+        -- relaxing of deriving mechanism to allow more types of deriving
+        DeriveAnyClass,
+        DeriveGeneric,
+        GeneralizedNewtypeDeriving,
+        StandaloneDeriving,
+        GADTs,
+        -- Convenience stuff
+        LambdaCase,
+        -- some HKT stuff
+        TypeOperators,
+        TypeApplications
+
     ghc-options: -fplugin=Effectful.Plugin
 
     hs-source-dirs: lib

From 681dba97fe4bf2cc8cf537337cb305d22d540f6f Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 2 Dec 2024 09:54:28 +0200
Subject: [PATCH 13/31] feat: some file structure refactoring

---
 .envrc => sky-node-gerbil/.envrc                       | 0
 Makefile => sky-node-gerbil/Makefile                   | 0
 all-sky.ss => sky-node-gerbil/all-sky.ss               | 0
 build.ss => sky-node-gerbil/build.ss                   | 0
 env.sh => sky-node-gerbil/env.sh                       | 0
 erasure-coding.ss => sky-node-gerbil/erasure-coding.ss | 0
 flake.lock => sky-node-gerbil/flake.lock               | 0
 flake.nix => sky-node-gerbil/flake.nix                 | 0
 gerbil.pkg => sky-node-gerbil/gerbil.pkg               | 0
 main.ss => sky-node-gerbil/main.ss                     | 0
 {pubsub => sky-node-gerbil/pubsub}/command.ss          | 0
 {pubsub => sky-node-gerbil/pubsub}/handler.ss          | 0
 {pubsub => sky-node-gerbil/pubsub}/lib.ss              | 0
 {pubsub => sky-node-gerbil/pubsub}/main.ss             | 0
 {pubsub => sky-node-gerbil/pubsub}/message.ss          | 0
 {pubsub => sky-node-gerbil/pubsub}/node.ss             | 0
 {t => sky-node-gerbil/t}/test-nodes.py                 | 0
 {t => sky-node-gerbil/t}/test.ss                       | 0
 unit-tests.ss => sky-node-gerbil/unit-tests.ss         | 0
 {sky-node => sky-node-haskell}/.envrc                  | 0
 {sky-node => sky-node-haskell}/CHANGELOG.md            | 0
 {sky-node => sky-node-haskell}/app/Main.hs             | 0
 {sky-node => sky-node-haskell}/config/default.yaml     | 0
 {sky-node => sky-node-haskell}/flake.lock              | 0
 {sky-node => sky-node-haskell}/flake.nix               | 0
 {sky-node => sky-node-haskell}/hie.yaml                | 0
 {sky-node => sky-node-haskell}/lib/App.hs              | 0
 {sky-node => sky-node-haskell}/lib/App/Env.hs          | 1 -
 {sky-node => sky-node-haskell}/lib/App/Error.hs        | 0
 {sky-node => sky-node-haskell}/lib/Command.hs          | 0
 {sky-node => sky-node-haskell}/lib/Config.hs           | 0
 {sky-node => sky-node-haskell}/lib/Peer.hs             | 0
 {sky-node => sky-node-haskell}/lib/Peer/Handlers.hs    | 0
 {sky-node => sky-node-haskell}/lib/Types.hs            | 0
 {sky-node => sky-node-haskell}/lib/Utils.hs            | 0
 {sky-node => sky-node-haskell}/sky-node.cabal          | 0
 36 files changed, 1 deletion(-)
 rename .envrc => sky-node-gerbil/.envrc (100%)
 rename Makefile => sky-node-gerbil/Makefile (100%)
 rename all-sky.ss => sky-node-gerbil/all-sky.ss (100%)
 rename build.ss => sky-node-gerbil/build.ss (100%)
 rename env.sh => sky-node-gerbil/env.sh (100%)
 rename erasure-coding.ss => sky-node-gerbil/erasure-coding.ss (100%)
 rename flake.lock => sky-node-gerbil/flake.lock (100%)
 rename flake.nix => sky-node-gerbil/flake.nix (100%)
 rename gerbil.pkg => sky-node-gerbil/gerbil.pkg (100%)
 rename main.ss => sky-node-gerbil/main.ss (100%)
 rename {pubsub => sky-node-gerbil/pubsub}/command.ss (100%)
 rename {pubsub => sky-node-gerbil/pubsub}/handler.ss (100%)
 rename {pubsub => sky-node-gerbil/pubsub}/lib.ss (100%)
 rename {pubsub => sky-node-gerbil/pubsub}/main.ss (100%)
 rename {pubsub => sky-node-gerbil/pubsub}/message.ss (100%)
 rename {pubsub => sky-node-gerbil/pubsub}/node.ss (100%)
 rename {t => sky-node-gerbil/t}/test-nodes.py (100%)
 rename {t => sky-node-gerbil/t}/test.ss (100%)
 rename unit-tests.ss => sky-node-gerbil/unit-tests.ss (100%)
 rename {sky-node => sky-node-haskell}/.envrc (100%)
 rename {sky-node => sky-node-haskell}/CHANGELOG.md (100%)
 rename {sky-node => sky-node-haskell}/app/Main.hs (100%)
 rename {sky-node => sky-node-haskell}/config/default.yaml (100%)
 rename {sky-node => sky-node-haskell}/flake.lock (100%)
 rename {sky-node => sky-node-haskell}/flake.nix (100%)
 rename {sky-node => sky-node-haskell}/hie.yaml (100%)
 rename {sky-node => sky-node-haskell}/lib/App.hs (100%)
 rename {sky-node => sky-node-haskell}/lib/App/Env.hs (90%)
 rename {sky-node => sky-node-haskell}/lib/App/Error.hs (100%)
 rename {sky-node => sky-node-haskell}/lib/Command.hs (100%)
 rename {sky-node => sky-node-haskell}/lib/Config.hs (100%)
 rename {sky-node => sky-node-haskell}/lib/Peer.hs (100%)
 rename {sky-node => sky-node-haskell}/lib/Peer/Handlers.hs (100%)
 rename {sky-node => sky-node-haskell}/lib/Types.hs (100%)
 rename {sky-node => sky-node-haskell}/lib/Utils.hs (100%)
 rename {sky-node => sky-node-haskell}/sky-node.cabal (100%)

diff --git a/.envrc b/sky-node-gerbil/.envrc
similarity index 100%
rename from .envrc
rename to sky-node-gerbil/.envrc
diff --git a/Makefile b/sky-node-gerbil/Makefile
similarity index 100%
rename from Makefile
rename to sky-node-gerbil/Makefile
diff --git a/all-sky.ss b/sky-node-gerbil/all-sky.ss
similarity index 100%
rename from all-sky.ss
rename to sky-node-gerbil/all-sky.ss
diff --git a/build.ss b/sky-node-gerbil/build.ss
similarity index 100%
rename from build.ss
rename to sky-node-gerbil/build.ss
diff --git a/env.sh b/sky-node-gerbil/env.sh
similarity index 100%
rename from env.sh
rename to sky-node-gerbil/env.sh
diff --git a/erasure-coding.ss b/sky-node-gerbil/erasure-coding.ss
similarity index 100%
rename from erasure-coding.ss
rename to sky-node-gerbil/erasure-coding.ss
diff --git a/flake.lock b/sky-node-gerbil/flake.lock
similarity index 100%
rename from flake.lock
rename to sky-node-gerbil/flake.lock
diff --git a/flake.nix b/sky-node-gerbil/flake.nix
similarity index 100%
rename from flake.nix
rename to sky-node-gerbil/flake.nix
diff --git a/gerbil.pkg b/sky-node-gerbil/gerbil.pkg
similarity index 100%
rename from gerbil.pkg
rename to sky-node-gerbil/gerbil.pkg
diff --git a/main.ss b/sky-node-gerbil/main.ss
similarity index 100%
rename from main.ss
rename to sky-node-gerbil/main.ss
diff --git a/pubsub/command.ss b/sky-node-gerbil/pubsub/command.ss
similarity index 100%
rename from pubsub/command.ss
rename to sky-node-gerbil/pubsub/command.ss
diff --git a/pubsub/handler.ss b/sky-node-gerbil/pubsub/handler.ss
similarity index 100%
rename from pubsub/handler.ss
rename to sky-node-gerbil/pubsub/handler.ss
diff --git a/pubsub/lib.ss b/sky-node-gerbil/pubsub/lib.ss
similarity index 100%
rename from pubsub/lib.ss
rename to sky-node-gerbil/pubsub/lib.ss
diff --git a/pubsub/main.ss b/sky-node-gerbil/pubsub/main.ss
similarity index 100%
rename from pubsub/main.ss
rename to sky-node-gerbil/pubsub/main.ss
diff --git a/pubsub/message.ss b/sky-node-gerbil/pubsub/message.ss
similarity index 100%
rename from pubsub/message.ss
rename to sky-node-gerbil/pubsub/message.ss
diff --git a/pubsub/node.ss b/sky-node-gerbil/pubsub/node.ss
similarity index 100%
rename from pubsub/node.ss
rename to sky-node-gerbil/pubsub/node.ss
diff --git a/t/test-nodes.py b/sky-node-gerbil/t/test-nodes.py
similarity index 100%
rename from t/test-nodes.py
rename to sky-node-gerbil/t/test-nodes.py
diff --git a/t/test.ss b/sky-node-gerbil/t/test.ss
similarity index 100%
rename from t/test.ss
rename to sky-node-gerbil/t/test.ss
diff --git a/unit-tests.ss b/sky-node-gerbil/unit-tests.ss
similarity index 100%
rename from unit-tests.ss
rename to sky-node-gerbil/unit-tests.ss
diff --git a/sky-node/.envrc b/sky-node-haskell/.envrc
similarity index 100%
rename from sky-node/.envrc
rename to sky-node-haskell/.envrc
diff --git a/sky-node/CHANGELOG.md b/sky-node-haskell/CHANGELOG.md
similarity index 100%
rename from sky-node/CHANGELOG.md
rename to sky-node-haskell/CHANGELOG.md
diff --git a/sky-node/app/Main.hs b/sky-node-haskell/app/Main.hs
similarity index 100%
rename from sky-node/app/Main.hs
rename to sky-node-haskell/app/Main.hs
diff --git a/sky-node/config/default.yaml b/sky-node-haskell/config/default.yaml
similarity index 100%
rename from sky-node/config/default.yaml
rename to sky-node-haskell/config/default.yaml
diff --git a/sky-node/flake.lock b/sky-node-haskell/flake.lock
similarity index 100%
rename from sky-node/flake.lock
rename to sky-node-haskell/flake.lock
diff --git a/sky-node/flake.nix b/sky-node-haskell/flake.nix
similarity index 100%
rename from sky-node/flake.nix
rename to sky-node-haskell/flake.nix
diff --git a/sky-node/hie.yaml b/sky-node-haskell/hie.yaml
similarity index 100%
rename from sky-node/hie.yaml
rename to sky-node-haskell/hie.yaml
diff --git a/sky-node/lib/App.hs b/sky-node-haskell/lib/App.hs
similarity index 100%
rename from sky-node/lib/App.hs
rename to sky-node-haskell/lib/App.hs
diff --git a/sky-node/lib/App/Env.hs b/sky-node-haskell/lib/App/Env.hs
similarity index 90%
rename from sky-node/lib/App/Env.hs
rename to sky-node-haskell/lib/App/Env.hs
index de81f8c..a4eaf68 100644
--- a/sky-node/lib/App/Env.hs
+++ b/sky-node-haskell/lib/App/Env.hs
@@ -1,5 +1,4 @@
 {-# LANGUAGE DataKinds #-}
--- while looks scary, it's needed because GHC can't guarantee HasField instances. Just be careful when defining them.
 {-# LANGUAGE UndecidableInstances #-}
 
 module App.Env
diff --git a/sky-node/lib/App/Error.hs b/sky-node-haskell/lib/App/Error.hs
similarity index 100%
rename from sky-node/lib/App/Error.hs
rename to sky-node-haskell/lib/App/Error.hs
diff --git a/sky-node/lib/Command.hs b/sky-node-haskell/lib/Command.hs
similarity index 100%
rename from sky-node/lib/Command.hs
rename to sky-node-haskell/lib/Command.hs
diff --git a/sky-node/lib/Config.hs b/sky-node-haskell/lib/Config.hs
similarity index 100%
rename from sky-node/lib/Config.hs
rename to sky-node-haskell/lib/Config.hs
diff --git a/sky-node/lib/Peer.hs b/sky-node-haskell/lib/Peer.hs
similarity index 100%
rename from sky-node/lib/Peer.hs
rename to sky-node-haskell/lib/Peer.hs
diff --git a/sky-node/lib/Peer/Handlers.hs b/sky-node-haskell/lib/Peer/Handlers.hs
similarity index 100%
rename from sky-node/lib/Peer/Handlers.hs
rename to sky-node-haskell/lib/Peer/Handlers.hs
diff --git a/sky-node/lib/Types.hs b/sky-node-haskell/lib/Types.hs
similarity index 100%
rename from sky-node/lib/Types.hs
rename to sky-node-haskell/lib/Types.hs
diff --git a/sky-node/lib/Utils.hs b/sky-node-haskell/lib/Utils.hs
similarity index 100%
rename from sky-node/lib/Utils.hs
rename to sky-node-haskell/lib/Utils.hs
diff --git a/sky-node/sky-node.cabal b/sky-node-haskell/sky-node.cabal
similarity index 100%
rename from sky-node/sky-node.cabal
rename to sky-node-haskell/sky-node.cabal

From 7d199bd08627d8960dc6ebaf5f0c999389387790 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 2 Dec 2024 11:50:28 +0200
Subject: [PATCH 14/31] feat: some package management

---
 sky-node-haskell/cabal.project  |  8 ++++
 sky-node-haskell/flake.lock     | 74 ++++++++++++++++++++++++++++++++-
 sky-node-haskell/flake.nix      |  7 +++-
 sky-node-haskell/sky-node.cabal |  1 +
 4 files changed, 88 insertions(+), 2 deletions(-)
 create mode 100644 sky-node-haskell/cabal.project

diff --git a/sky-node-haskell/cabal.project b/sky-node-haskell/cabal.project
new file mode 100644
index 0000000..8065894
--- /dev/null
+++ b/sky-node-haskell/cabal.project
@@ -0,0 +1,8 @@
+packages: ./*.cabal
+documentation: True
+
+source-repository-package
+	type: git
+	location: https://github.com/SkyProtocol-org/skyprotocol/
+	tag: 6be2686e0f53b6f2ba31e557090bd5bf4f56289a
+	subdir: merkle-patricia-trie
diff --git a/sky-node-haskell/flake.lock b/sky-node-haskell/flake.lock
index d7d724e..ef4cf99 100644
--- a/sky-node-haskell/flake.lock
+++ b/sky-node-haskell/flake.lock
@@ -18,7 +18,63 @@
         "type": "github"
       }
     },
+    "flake-utils_2": {
+      "inputs": {
+        "systems": "systems_2"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "merkle-trie": {
+      "inputs": {
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "dir": "merkle-patricia-trie",
+        "lastModified": 1733057365,
+        "narHash": "sha256-zv4Ezyes0jSmCTs2EZUz8Ckotwaj+kbElmBdcXxU+xo=",
+        "owner": "Skyprotocol-org",
+        "repo": "skyprotocol",
+        "rev": "6be2686e0f53b6f2ba31e557090bd5bf4f56289a",
+        "type": "github"
+      },
+      "original": {
+        "dir": "merkle-patricia-trie",
+        "owner": "Skyprotocol-org",
+        "ref": "issue-22/haskell-implementation-of-sky-node",
+        "repo": "skyprotocol",
+        "type": "github"
+      }
+    },
     "nixpkgs": {
+      "locked": {
+        "lastModified": 1732758367,
+        "narHash": "sha256-RzaI1RO0UXqLjydtz3GAXSTzHkpb/lLD1JD8a0W4Wpo=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "fa42b5a5f401aab8a32bd33c9a4de0738180dc59",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "ref": "nixos-unstable",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "nixpkgs_2": {
       "locked": {
         "lastModified": 1731319897,
         "narHash": "sha256-PbABj4tnbWFMfBp6OcUK5iGy1QY+/Z96ZcLpooIbuEI=",
@@ -37,7 +93,8 @@
     "root": {
       "inputs": {
         "flake-utils": "flake-utils",
-        "nixpkgs": "nixpkgs"
+        "merkle-trie": "merkle-trie",
+        "nixpkgs": "nixpkgs_2"
       }
     },
     "systems": {
@@ -54,6 +111,21 @@
         "repo": "default",
         "type": "github"
       }
+    },
+    "systems_2": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
     }
   },
   "root": "root",
diff --git a/sky-node-haskell/flake.nix b/sky-node-haskell/flake.nix
index b08fd04..9309932 100644
--- a/sky-node-haskell/flake.nix
+++ b/sky-node-haskell/flake.nix
@@ -2,12 +2,16 @@
   inputs = {
     flake-utils.url = "github:numtide/flake-utils";
     nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
+    merkle-trie = {
+      url = "github:Skyprotocol-org/skyprotocol/issue-22/haskell-implementation-of-sky-node?dir=merkle-patricia-trie";
+    };
   };
-  outputs = { self, nixpkgs, flake-utils, ... }: flake-utils.lib.eachDefaultSystem (system:
+  outputs = { self, nixpkgs, flake-utils, merkle-trie, ... }: flake-utils.lib.eachDefaultSystem (system:
     let
       pkgs = import nixpkgs { inherit system; };
       overlay = final: prev: {
         sky-node = final.callCabal2nix "sky-node" ./. { };
+        merkle-trie = final.callCabal2nix "merkle-patricia-trie" merkle-trie.src { };
       };
       haskPkgs = pkgs.haskellPackages.extend overlay;
     in
@@ -15,6 +19,7 @@
       devShells.default = haskPkgs.shellFor {
         packages = p: [
           p.sky-node
+          p.merkle-patricia-trie
         ];
         nativeBuildInputs = with haskPkgs; [
           cabal-install
diff --git a/sky-node-haskell/sky-node.cabal b/sky-node-haskell/sky-node.cabal
index b47a440..893530e 100644
--- a/sky-node-haskell/sky-node.cabal
+++ b/sky-node-haskell/sky-node.cabal
@@ -91,6 +91,7 @@ library sky-lib
         exceptions,
         stm,
         data-default,
+        merkle-patricia-trie
 
     default-extensions:
         -- enables {..} syntax for records

From 180488edf2b43323f0a426afc57d132fe0bf0123 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 2 Dec 2024 13:55:36 +0200
Subject: [PATCH 15/31] fix: modified flake to see the merkle-patricia-trie as
 dependency of sky-node

---
 sky-node-haskell/cabal.project | 8 --------
 sky-node-haskell/flake.nix     | 5 ++---
 2 files changed, 2 insertions(+), 11 deletions(-)
 delete mode 100644 sky-node-haskell/cabal.project

diff --git a/sky-node-haskell/cabal.project b/sky-node-haskell/cabal.project
deleted file mode 100644
index 8065894..0000000
--- a/sky-node-haskell/cabal.project
+++ /dev/null
@@ -1,8 +0,0 @@
-packages: ./*.cabal
-documentation: True
-
-source-repository-package
-	type: git
-	location: https://github.com/SkyProtocol-org/skyprotocol/
-	tag: 6be2686e0f53b6f2ba31e557090bd5bf4f56289a
-	subdir: merkle-patricia-trie
diff --git a/sky-node-haskell/flake.nix b/sky-node-haskell/flake.nix
index 9309932..12780ac 100644
--- a/sky-node-haskell/flake.nix
+++ b/sky-node-haskell/flake.nix
@@ -10,8 +10,8 @@
     let
       pkgs = import nixpkgs { inherit system; };
       overlay = final: prev: {
-        sky-node = final.callCabal2nix "sky-node" ./. { };
-        merkle-trie = final.callCabal2nix "merkle-patricia-trie" merkle-trie.src { };
+        merkle-patricia-trie = prev.callCabal2nix "merkle-patricia-trie" "${merkle-trie}" { };
+        sky-node = prev.callCabal2nix "sky-node" ./. { };
       };
       haskPkgs = pkgs.haskellPackages.extend overlay;
     in
@@ -19,7 +19,6 @@
       devShells.default = haskPkgs.shellFor {
         packages = p: [
           p.sky-node
-          p.merkle-patricia-trie
         ];
         nativeBuildInputs = with haskPkgs; [
           cabal-install

From ef33aea2fb86fcac78d03ad0bf936eadadd3ee4a Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Wed, 4 Dec 2024 18:30:22 +0200
Subject: [PATCH 16/31] fix: small refactoring

---
 sky-node-haskell/lib/App.hs           | 2 +-
 sky-node-haskell/lib/Peer/Handlers.hs | 4 ++--
 sky-node-haskell/lib/Types.hs         | 8 ++++----
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/sky-node-haskell/lib/App.hs b/sky-node-haskell/lib/App.hs
index d3a54d6..3ee15d1 100644
--- a/sky-node-haskell/lib/App.hs
+++ b/sky-node-haskell/lib/App.hs
@@ -35,7 +35,7 @@ withSocketServer handler = do
     forever $ do
       (conn, addr) <- liftIO $ S.accept sock
       logInfo_ $ "Connection accepted from " <> pack (show addr)
-      -- using withAsync ensures proper thread handling in face of execption
+      -- using withAsync ensures proper thread handling in the face of execption
       withAsync (handler conn) $ \_ -> pure ()
 
 setupServerSocket :: Eff AppEffects S.Socket
diff --git a/sky-node-haskell/lib/Peer/Handlers.hs b/sky-node-haskell/lib/Peer/Handlers.hs
index 530b971..abbbefb 100644
--- a/sky-node-haskell/lib/Peer/Handlers.hs
+++ b/sky-node-haskell/lib/Peer/Handlers.hs
@@ -37,7 +37,7 @@ runPeerHandlerIO = interpret $ \_ -> \case
     topics <- askFieldS @Topics
     adapt $ atomically $ do
       tpcs <- readTVar topics
-      let tpcs' = IntMap.adjust (\tpc -> tpc {topicMessages = IntMap.insert (IntMap.size tpc.topicMessages) bData tpc.topicMessages}) topicId tpcs
+      let tpcs' = IntMap.adjust (\tpc -> tpc {messages = IntMap.insert (IntMap.size tpc.messages) bData tpc.messages}) id tpcs
       writeTVar topics tpcs'
     pure $ makeBlockCertificate bData
   GetTopics mtId -> adapt $ undefined
@@ -46,7 +46,7 @@ runPeerHandlerIO = interpret $ \_ -> \case
     adapt $ do
       maybeMeta <- atomically $ do
         tpcs <- readTVar topics
-        let mTpc = (.topicMeta) <$> (IntMap.!?) tpcs topicId
+        let mTpc = (.metadata) <$> (IntMap.!?) tpcs id
         pure mTpc
       case maybeMeta of
         Just meta -> pure meta
diff --git a/sky-node-haskell/lib/Types.hs b/sky-node-haskell/lib/Types.hs
index 21fb8e4..0c5c15e 100644
--- a/sky-node-haskell/lib/Types.hs
+++ b/sky-node-haskell/lib/Types.hs
@@ -3,13 +3,13 @@ module Types where
 import qualified Data.ByteString as BS
 import Data.IntMap.Strict (IntMap)
 
-newtype TopicId = TopicId {topicId :: Int} deriving (Show, Eq)
+newtype TopicId = TopicId {id :: Int} deriving (Show, Eq)
 
-newtype TopicMetaData = TopicMetaData {topicMetaId :: TopicId} deriving (Show, Eq)
+newtype TopicMetaData = TopicMetaData {id :: TopicId} deriving (Show, Eq)
 
 data Topic = Topic
-  { topicMeta :: TopicMetaData,
-    topicMessages :: IntMap BlockData
+  { metadata :: TopicMetaData,
+    messages :: IntMap BlockData
   }
   deriving (Show, Eq)
 

From 268082dfa2a2e14dd91b5d29383ec322a1556575 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Thu, 5 Dec 2024 00:25:06 +0200
Subject: [PATCH 17/31] chore: some refactoring of effects and project
 structure

---
 sky-node-haskell/lib/App.hs                   | 20 +++++++----
 sky-node-haskell/lib/Command.hs               |  0
 sky-node-haskell/lib/Effect.hs                |  1 +
 .../{Peer/Handlers.hs => Effect/PeerEff.hs}   | 35 +++++++++----------
 sky-node-haskell/lib/Effect/SocketEff.hs      | 23 ++++++++++++
 sky-node-haskell/lib/Peer.hs                  |  6 ++--
 sky-node-haskell/lib/Types.hs                 | 28 ++++++++++-----
 sky-node-haskell/sky-node.cabal               |  6 +++-
 8 files changed, 84 insertions(+), 35 deletions(-)
 delete mode 100644 sky-node-haskell/lib/Command.hs
 create mode 100644 sky-node-haskell/lib/Effect.hs
 rename sky-node-haskell/lib/{Peer/Handlers.hs => Effect/PeerEff.hs} (57%)
 create mode 100644 sky-node-haskell/lib/Effect/SocketEff.hs

diff --git a/sky-node-haskell/lib/App.hs b/sky-node-haskell/lib/App.hs
index 3ee15d1..7a67fa6 100644
--- a/sky-node-haskell/lib/App.hs
+++ b/sky-node-haskell/lib/App.hs
@@ -6,25 +6,32 @@ import App.Env
 import Config (AppConfig (..))
 import Control.Monad (forever)
 import Data.Text (pack, unpack)
+import Effect.PeerEff
 import Effectful
 import Effectful.Concurrent
 import Effectful.Concurrent.Async
+import Effectful.Error.Static (Error, runError)
 import Effectful.Log
 import Effectful.Reader.Static
 import qualified Network.Socket as S
 import UnliftIO.Exception (bracket)
 import Utils
 
-type AppEffects = '[Reader AppEnv, Concurrent, Log, IOE]
+type AppEffects = '[Reader AppEnv, Concurrent, Error String, Log, IOE]
 
 initApp :: AppConfig -> Logger -> IO ()
 initApp config logger = do
   env <- initAppEnv config
   runEff $ do
-    runLog "main" logger defaultLogLevel $
-      runConcurrent $
-        runReader env $ do
-          withSocketServer handlePeer
+    runLog "main" logger defaultLogLevel $ do
+      -- we don't want handlers to interrupt the node with it's exceptions
+      eRes <- runError $
+        runConcurrent $
+          runReader env $ do
+            withSocketServer handlePeer
+      case eRes of
+        Left err -> logAttention_ $ pack $ show err
+        Right res -> logInfo_ $ pack $ show res
 
 withSocketServer :: (S.Socket -> Eff AppEffects ()) -> Eff AppEffects ()
 withSocketServer handler = do
@@ -52,4 +59,5 @@ setupServerSocket = do
   pure sock
 
 handlePeer :: S.Socket -> Eff AppEffects ()
-handlePeer sock = undefined
+handlePeer sock = runPeerEffIO sock $ do
+  undefined
diff --git a/sky-node-haskell/lib/Command.hs b/sky-node-haskell/lib/Command.hs
deleted file mode 100644
index e69de29..0000000
diff --git a/sky-node-haskell/lib/Effect.hs b/sky-node-haskell/lib/Effect.hs
new file mode 100644
index 0000000..fca483c
--- /dev/null
+++ b/sky-node-haskell/lib/Effect.hs
@@ -0,0 +1 @@
+module Effect where
diff --git a/sky-node-haskell/lib/Peer/Handlers.hs b/sky-node-haskell/lib/Effect/PeerEff.hs
similarity index 57%
rename from sky-node-haskell/lib/Peer/Handlers.hs
rename to sky-node-haskell/lib/Effect/PeerEff.hs
index abbbefb..616f787 100644
--- a/sky-node-haskell/lib/Peer/Handlers.hs
+++ b/sky-node-haskell/lib/Effect/PeerEff.hs
@@ -1,11 +1,9 @@
 {-# LANGUAGE DataKinds #-}
-{-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE TypeOperators #-}
 
-module Peer.Handlers where
+module Effect.PeerEff where
 
 import App.Env (Topics)
 import Control.Concurrent.STM (atomically, readTVar, writeTVar)
@@ -16,37 +14,38 @@ import Effectful.Dispatch.Dynamic
 import Effectful.Error.Static
 import Effectful.Reader.Static (Reader)
 import Effectful.TH (makeEffect)
+import qualified Network.Socket as S
 import Types
 import UnliftIO (catchIO, throwString)
 import Utils
 
-data PeerHandler :: Effect where
-  PublishBlock :: TopicId -> BlockData -> PeerHandler m (Certificate Block)
-  GetTopics :: Maybe TopicId -> PeerHandler m [TopicId]
-  DescribeTopic :: TopicId -> PeerHandler m TopicMetaData
-  PollTopic :: TopicId -> PeerHandler m (Word64, Certificate Block)
-  GetTopicBlockCertificate :: TopicId -> Word64 -> PeerHandler m (Certificate Block)
-  ReadTopic :: TopicId -> Word64 -> PeerHandler m [BlockData]
+data PeerEff :: Effect where
+  PublishBlock :: TopicId -> BlockData -> PeerEff m Certificate
+  GetTopics :: Maybe TopicId -> PeerEff m [TopicId]
+  DescribeTopic :: TopicId -> PeerEff m TopicMetaData
+  PollTopic :: TopicId -> PeerEff m (Word64, Certificate)
+  GetTopicBlockCertificate :: TopicId -> Word64 -> PeerEff m Certificate
+  ReadTopic :: TopicId -> Word64 -> PeerEff m [BlockData]
 
 -- TODO replace with makeEffect_ and add docs for type sigs, or add docs to the GADT
-makeEffect ''PeerHandler
+makeEffect ''PeerEff
 
-runPeerHandlerIO :: (Has Topics env, IOE :> es, Error String :> es, Reader env :> es) => Eff (PeerHandler : es) a -> Eff es a
-runPeerHandlerIO = interpret $ \_ -> \case
-  PublishBlock TopicId {..} bData -> do
+runPeerEffIO :: (Has Topics env, IOE :> es, Error String :> es, Reader env :> es) => S.Socket -> Eff (PeerEff : es) a -> Eff es a
+runPeerEffIO sock = interpret $ \_ -> \case
+  PublishBlock tId bData -> do
     topics <- askFieldS @Topics
     adapt $ atomically $ do
       tpcs <- readTVar topics
-      let tpcs' = IntMap.adjust (\tpc -> tpc {messages = IntMap.insert (IntMap.size tpc.messages) bData tpc.messages}) id tpcs
+      let tpcs' = IntMap.adjust (\tpc -> tpc {messages = IntMap.insert (IntMap.size tpc.messages) bData tpc.messages}) tId.id tpcs
       writeTVar topics tpcs'
-    pure $ makeBlockCertificate bData
+    pure $ Certificate ""
   GetTopics mtId -> adapt $ undefined
-  DescribeTopic TopicId {..} -> do
+  DescribeTopic tId -> do
     topics <- askFieldS @Topics
     adapt $ do
       maybeMeta <- atomically $ do
         tpcs <- readTVar topics
-        let mTpc = (.metadata) <$> (IntMap.!?) tpcs id
+        let mTpc = (.metadata) <$> (IntMap.!?) tpcs tId.id
         pure mTpc
       case maybeMeta of
         Just meta -> pure meta
diff --git a/sky-node-haskell/lib/Effect/SocketEff.hs b/sky-node-haskell/lib/Effect/SocketEff.hs
new file mode 100644
index 0000000..6ec66fb
--- /dev/null
+++ b/sky-node-haskell/lib/Effect/SocketEff.hs
@@ -0,0 +1,23 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Effect.SocketEff where
+
+import Effectful
+import Effectful.Dispatch.Dynamic (interpret)
+import Effectful.TH (makeEffect)
+import qualified Network.Socket as S
+
+data SocketEff :: Effect where
+  RecvCommand :: SocketEff m ()
+  SendCommand :: SocketEff m ()
+
+-- TODO replace with makeEffect_ and add docs for type sigs, or add docs to the GADT
+makeEffect ''SocketEff
+
+runSocketEffIO :: (IOE :> es) => S.Socket -> Eff (SocketEff : es) a -> Eff es a
+runSocketEffIO sock = interpret $ \_ -> \case
+  RecvCommand -> undefined
+  SendCommand -> undefined
diff --git a/sky-node-haskell/lib/Peer.hs b/sky-node-haskell/lib/Peer.hs
index 1e13491..8b039c9 100644
--- a/sky-node-haskell/lib/Peer.hs
+++ b/sky-node-haskell/lib/Peer.hs
@@ -3,9 +3,11 @@ module Peer
   )
 where
 
+import qualified Data.ByteString as BS
+
 -- | Aux data structure to keep info about 'Peer'.
 newtype Peer = Peer
   { -- | Id of the 'Peer'
-    id :: String
+    id :: BS.ByteString
   }
-  deriving (Show)
+  deriving (Show, Eq)
diff --git a/sky-node-haskell/lib/Types.hs b/sky-node-haskell/lib/Types.hs
index 0c5c15e..e86885d 100644
--- a/sky-node-haskell/lib/Types.hs
+++ b/sky-node-haskell/lib/Types.hs
@@ -1,24 +1,36 @@
 module Types where
 
+import Data.Binary
 import qualified Data.ByteString as BS
 import Data.IntMap.Strict (IntMap)
+import GHC.Generics
 
-newtype TopicId = TopicId {id :: Int} deriving (Show, Eq)
+newtype TopicId = TopicId {id :: Int}
+  deriving stock (Show, Eq, Generic)
 
-newtype TopicMetaData = TopicMetaData {id :: TopicId} deriving (Show, Eq)
+instance Binary TopicId
+
+newtype TopicMetaData = TopicMetaData {id :: TopicId}
+  deriving stock (Show, Eq, Generic)
+
+instance Binary TopicMetaData
 
 data Topic = Topic
   { metadata :: TopicMetaData,
     messages :: IntMap BlockData
   }
-  deriving (Show, Eq)
+  deriving stock (Show, Eq, Generic)
+
+instance Binary Topic
 
 data Block
 
-newtype BlockData = BlockData {blockData :: BS.ByteString} deriving (Show, Eq)
+newtype BlockData = BlockData {blockData :: BS.ByteString}
+  deriving stock (Show, Eq, Generic)
+
+instance Binary BlockData
 
-newtype Certificate a = Certificate {cert :: BS.ByteString} deriving (Show, Eq)
+newtype Certificate = Certificate {cert :: BS.ByteString}
+  deriving stock (Show, Eq, Generic)
 
--- TODO for now it's a stub
-makeBlockCertificate :: BlockData -> Certificate Block
-makeBlockCertificate = undefined
+instance Binary Certificate
diff --git a/sky-node-haskell/sky-node.cabal b/sky-node-haskell/sky-node.cabal
index 893530e..eddcd50 100644
--- a/sky-node-haskell/sky-node.cabal
+++ b/sky-node-haskell/sky-node.cabal
@@ -70,11 +70,14 @@ library sky-lib
         App.Env
         Config
         Peer
-        Peer.Handlers
+        Effect
+        Effect.PeerEff
+        Effect.SocketEff
         Types
         Utils
 
     build-depends:
+        binary,
         network,
         effectful,
         effectful-core,
@@ -117,6 +120,7 @@ library sky-lib
         DeriveGeneric,
         GeneralizedNewtypeDeriving,
         StandaloneDeriving,
+        DerivingStrategies,
         GADTs,
         -- Convenience stuff
         LambdaCase,

From 8ce1c89167d31410ac1da583b160bccab59b59b4 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Wed, 11 Dec 2024 00:03:13 +0200
Subject: [PATCH 18/31] feat: trie implementation

---
 merkle-patricia-trie/lib/Data/Trie.hs         | 75 +++++++++++++++++++
 .../merkle-patricia-trie.cabal                |  7 +-
 2 files changed, 80 insertions(+), 2 deletions(-)
 create mode 100644 merkle-patricia-trie/lib/Data/Trie.hs

diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
new file mode 100644
index 0000000..aba35a7
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -0,0 +1,75 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Data.Trie (Trie (..), lookup, insert) where
+
+import Data.Bits
+import Prelude (Bool, Enum, Eq, Integral, Maybe (..), Num, Ord, Real, not, otherwise, undefined, ($), (-), (.), (==))
+
+data Trie k a
+  = Empty
+  | Leaf k a
+  | Branch k k (Trie k a) (Trie k a)
+
+-- the idea is that you create a datatype for key, which
+class (Bits k, Num k, Integral k) => TrieKey k where
+  mask :: k -> k -> k
+  branchingBit :: k -> k -> k
+
+newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
+
+newtype BEK a = BigEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
+
+instance (Bits k, Num k, Integral k) => TrieKey (LEK k) where
+  mask (LittleEndian a) (LittleEndian b) = LittleEndian $ a .&. (b - 1)
+  branchingBit (LittleEndian a) (LittleEndian b) = LittleEndian . lowestBit $ a `xor` b
+    where
+      lowestBit x = x .&. complement x
+
+instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
+  mask (BigEndian a) (BigEndian b) = BigEndian $ (a .|. (b - 1)) .&. complement b
+
+  -- TODO figure out how to properly implement branchingBit for BigEndian
+  -- since it needs more args than just a and b(it also needs prefixes)
+  branchingBit _ _ = undefined
+
+lookup :: (TrieKey k) => k -> Trie k a -> Maybe a
+lookup _ Empty = Nothing
+lookup k (Leaf key a) = if k == key then Just a else Nothing
+lookup k (Branch key prefix t1 t2)
+  | not (matchPrefix k prefix key) = Nothing
+  | zeroBit k key = lookup k t1
+  | otherwise = lookup k t2
+
+matchPrefix :: (TrieKey k) => k -> k -> k -> Bool
+matchPrefix k prefix key = mask k key == prefix
+
+zeroBit :: (TrieKey k) => k -> k -> Bool
+zeroBit a b = (a .&. b) == 0
+
+join :: (TrieKey k) => k -> Trie k a -> k -> Trie k a -> Trie k a
+join p1 t1 p2 t2 =
+  if zeroBit p1 commonPrefix
+    then Branch (mask p1 commonPrefix) commonPrefix t1 t2
+    else Branch (mask p1 commonPrefix) commonPrefix t2 t1
+  where
+    commonPrefix = branchingBit p1 p2
+
+-- maybe conflict func must be a Monoid?
+insert :: forall k v. (TrieKey k) => (v -> v -> v) -> k -> v -> Trie k v -> Trie k v
+insert conflict k v t = go t
+  where
+    go :: Trie k v -> Trie k v
+    go Empty = Leaf k v
+    go (Leaf key val) =
+      if k == key
+        then Leaf k $ conflict v val
+        else join k (Leaf k v) key t
+    go (Branch key prefix t1 t2) =
+      if matchPrefix k prefix key
+        then
+          if zeroBit k key
+            then Branch key prefix (go t1) t2
+            else Branch key prefix t1 (go t2)
+        else join k (Leaf k v) prefix t
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index ed41911..2127795 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -57,13 +57,16 @@ library
     import:           warnings
 
     -- Modules exported by the library.
-    exposed-modules:  Data.MerkleTrie
+    exposed-modules:
+        Data.MerkleTrie
+        Data.Trie
 
     -- Modules included in this library but not exported.
     -- other-modules:
 
     -- LANGUAGE extensions used by modules in this package.
-    -- other-extensions:
+    other-extensions:
+        NoImplicitPrelude
 
     -- Other library packages from which modules are imported.
     build-depends:

From dd4c51440aa9481b0c8bea9b7687ea798a7f79e7 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Wed, 11 Dec 2024 16:38:29 +0200
Subject: [PATCH 19/31] feat: small additions

---
 merkle-patricia-trie/lib/Data/Trie.hs | 29 ++++++++++++++-------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index aba35a7..18fc225 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,20 +1,22 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeFamilies #-}
 
 module Data.Trie (Trie (..), lookup, insert) where
 
 import Data.Bits
-import Prelude (Bool, Enum, Eq, Integral, Maybe (..), Num, Ord, Real, not, otherwise, undefined, ($), (-), (.), (==))
+import Data.Kind (Type)
+import Prelude hiding (lookup)
 
-data Trie k a
+data Trie (f :: Type -> Type) k a
   = Empty
-  | Leaf k a
-  | Branch k k (Trie k a) (Trie k a)
+  | Leaf k (f a)
+  | Branch k k (Trie f k a) (Trie f k a)
 
--- the idea is that you create a datatype for key, which
 class (Bits k, Num k, Integral k) => TrieKey k where
+  -- | Discards bits before/after(depending on endianess) branching bit
   mask :: k -> k -> k
+
+  -- | Leaves only the branching bit(so key 0b1001101 becomes 0b0001000, if 4th bit was branching)
   branchingBit :: k -> k -> k
 
 newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
@@ -34,7 +36,7 @@ instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
   -- since it needs more args than just a and b(it also needs prefixes)
   branchingBit _ _ = undefined
 
-lookup :: (TrieKey k) => k -> Trie k a -> Maybe a
+lookup :: (TrieKey k, Functor f) => k -> Trie f k a -> Maybe (f a)
 lookup _ Empty = Nothing
 lookup k (Leaf key a) = if k == key then Just a else Nothing
 lookup k (Branch key prefix t1 t2)
@@ -48,7 +50,7 @@ matchPrefix k prefix key = mask k key == prefix
 zeroBit :: (TrieKey k) => k -> k -> Bool
 zeroBit a b = (a .&. b) == 0
 
-join :: (TrieKey k) => k -> Trie k a -> k -> Trie k a -> Trie k a
+join :: (TrieKey k) => k -> Trie f k a -> k -> Trie f k a -> Trie f k a
 join p1 t1 p2 t2 =
   if zeroBit p1 commonPrefix
     then Branch (mask p1 commonPrefix) commonPrefix t1 t2
@@ -57,19 +59,18 @@ join p1 t1 p2 t2 =
     commonPrefix = branchingBit p1 p2
 
 -- maybe conflict func must be a Monoid?
-insert :: forall k v. (TrieKey k) => (v -> v -> v) -> k -> v -> Trie k v -> Trie k v
+insert :: (TrieKey k, Functor f, Applicative f) => (v -> v -> v) -> k -> v -> Trie f k v -> Trie f k v
 insert conflict k v t = go t
   where
-    go :: Trie k v -> Trie k v
-    go Empty = Leaf k v
+    go Empty = Leaf k $ pure v
     go (Leaf key val) =
       if k == key
-        then Leaf k $ conflict v val
-        else join k (Leaf k v) key t
+        then Leaf k $ conflict v <$> val
+        else join k (Leaf k $ pure v) key t
     go (Branch key prefix t1 t2) =
       if matchPrefix k prefix key
         then
           if zeroBit k key
             then Branch key prefix (go t1) t2
             else Branch key prefix t1 (go t2)
-        else join k (Leaf k v) prefix t
+        else join k (Leaf k $ pure v) prefix t

From e227be5816e7f9627def463ec540acb1f2756b89 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Thu, 12 Dec 2024 16:58:02 +0200
Subject: [PATCH 20/31] feat: initial zipper impl

---
 merkle-patricia-trie/lib/Data/Trie.hs         | 59 ++++---------------
 .../lib/Data/Trie/Internal.hs                 | 52 ++++++++++++++++
 merkle-patricia-trie/lib/Data/Trie/Zipper.hs  | 56 ++++++++++++++++++
 .../merkle-patricia-trie.cabal                |  2 +
 4 files changed, 121 insertions(+), 48 deletions(-)
 create mode 100644 merkle-patricia-trie/lib/Data/Trie/Internal.hs
 create mode 100644 merkle-patricia-trie/lib/Data/Trie/Zipper.hs

diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index 18fc225..cc638b5 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,41 +1,12 @@
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.Trie (Trie (..), lookup, insert) where
+module Data.Trie (Trie (..), TrieKey, lookup, insert) where
 
 import Data.Bits
 import Data.Kind (Type)
+import Data.Trie.Internal
 import Prelude hiding (lookup)
 
-data Trie (f :: Type -> Type) k a
-  = Empty
-  | Leaf k (f a)
-  | Branch k k (Trie f k a) (Trie f k a)
-
-class (Bits k, Num k, Integral k) => TrieKey k where
-  -- | Discards bits before/after(depending on endianess) branching bit
-  mask :: k -> k -> k
-
-  -- | Leaves only the branching bit(so key 0b1001101 becomes 0b0001000, if 4th bit was branching)
-  branchingBit :: k -> k -> k
-
-newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
-
-newtype BEK a = BigEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
-
-instance (Bits k, Num k, Integral k) => TrieKey (LEK k) where
-  mask (LittleEndian a) (LittleEndian b) = LittleEndian $ a .&. (b - 1)
-  branchingBit (LittleEndian a) (LittleEndian b) = LittleEndian . lowestBit $ a `xor` b
-    where
-      lowestBit x = x .&. complement x
-
-instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
-  mask (BigEndian a) (BigEndian b) = BigEndian $ (a .|. (b - 1)) .&. complement b
-
-  -- TODO figure out how to properly implement branchingBit for BigEndian
-  -- since it needs more args than just a and b(it also needs prefixes)
-  branchingBit _ _ = undefined
-
 lookup :: (TrieKey k, Functor f) => k -> Trie f k a -> Maybe (f a)
 lookup _ Empty = Nothing
 lookup k (Leaf key a) = if k == key then Just a else Nothing
@@ -44,28 +15,14 @@ lookup k (Branch key prefix t1 t2)
   | zeroBit k key = lookup k t1
   | otherwise = lookup k t2
 
-matchPrefix :: (TrieKey k) => k -> k -> k -> Bool
-matchPrefix k prefix key = mask k key == prefix
-
-zeroBit :: (TrieKey k) => k -> k -> Bool
-zeroBit a b = (a .&. b) == 0
-
-join :: (TrieKey k) => k -> Trie f k a -> k -> Trie f k a -> Trie f k a
-join p1 t1 p2 t2 =
-  if zeroBit p1 commonPrefix
-    then Branch (mask p1 commonPrefix) commonPrefix t1 t2
-    else Branch (mask p1 commonPrefix) commonPrefix t2 t1
-  where
-    commonPrefix = branchingBit p1 p2
-
 -- maybe conflict func must be a Monoid?
-insert :: (TrieKey k, Functor f, Applicative f) => (v -> v -> v) -> k -> v -> Trie f k v -> Trie f k v
-insert conflict k v t = go t
+insertWith :: (TrieKey k, Functor f, Applicative f) => (v -> v -> v) -> k -> v -> Trie f k v -> Trie f k v
+insertWith resolve k v t = go t
   where
     go Empty = Leaf k $ pure v
     go (Leaf key val) =
       if k == key
-        then Leaf k $ conflict v <$> val
+        then Leaf k $ resolve v <$> val
         else join k (Leaf k $ pure v) key t
     go (Branch key prefix t1 t2) =
       if matchPrefix k prefix key
@@ -74,3 +31,9 @@ insert conflict k v t = go t
             then Branch key prefix (go t1) t2
             else Branch key prefix t1 (go t2)
         else join k (Leaf k $ pure v) prefix t
+
+insert :: (TrieKey k, Functor f, Applicative f) => k -> v -> Trie f k v -> Trie f k v
+insert = insertWith const
+
+empty :: (TrieKey k, Functor f) => Trie f k v
+empty = Empty
diff --git a/merkle-patricia-trie/lib/Data/Trie/Internal.hs b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
new file mode 100644
index 0000000..c1e6cd7
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
@@ -0,0 +1,52 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module Data.Trie.Internal where
+
+import Data.Bits
+import Data.Kind (Type)
+import Prelude hiding (lookup)
+
+data Trie (f :: Type -> Type) k a
+  = Empty
+  | Leaf k (f a)
+  | Branch k k (Trie f k a) (Trie f k a)
+
+class (Bits k, Num k, Integral k) => TrieKey k where
+  -- | Discards bits before/after(depending on endianess) branching bit
+  mask :: k -> k -> k
+
+  -- | Leaves only the branching bit(so key 0b1001101 becomes 0b0001000, if 4th bit was branching)
+  branchingBit :: k -> k -> k
+
+newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
+
+newtype BEK a = BigEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
+
+instance (Bits k, Num k, Integral k) => TrieKey (LEK k) where
+  mask (LittleEndian a) (LittleEndian b) = LittleEndian $ a .&. (b - 1)
+  branchingBit (LittleEndian a) (LittleEndian b) = LittleEndian . lowestBit $ a `xor` b
+    where
+      lowestBit x = x .&. complement x
+
+instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
+  mask (BigEndian a) (BigEndian b) = BigEndian $ (a .|. (b - 1)) .&. complement b
+
+  -- TODO figure out how to properly implement branchingBit for BigEndian
+  -- since it needs more args than just a and b(it also needs prefixes)
+  branchingBit _ _ = undefined
+
+matchPrefix :: (TrieKey k) => k -> k -> k -> Bool
+matchPrefix k prefix key = mask k key == prefix
+
+zeroBit :: (TrieKey k) => k -> k -> Bool
+zeroBit a b = (a .&. b) == 0
+
+join :: (TrieKey k) => k -> Trie f k a -> k -> Trie f k a -> Trie f k a
+join p1 t1 p2 t2 =
+  if zeroBit p1 commonPrefix
+    then Branch (mask p1 commonPrefix) commonPrefix t1 t2
+    else Branch (mask p1 commonPrefix) commonPrefix t2 t1
+  where
+    commonPrefix = branchingBit p1 p2
diff --git a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
new file mode 100644
index 0000000..0edd40e
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
@@ -0,0 +1,56 @@
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE ViewPatterns #-}
+
+module Data.Trie.Zipper (TrieZipper, left, right, up, lookupZipper, zipper) where
+
+-- import Control.Monad.Identity (Identity (..))
+import Data.Kind (Type)
+import Data.Trie.Internal
+
+data Cxt (f :: Type -> Type) k a
+  = Top
+  | L k k (Cxt f k a) (Trie f k a)
+  | R k k (Trie f k a) (Cxt f k a)
+
+newtype TrieZipper (f :: Type -> Type) k a = TrieZipper {getZipper :: (Trie f k a, Cxt f k a)}
+
+zipper :: (TrieKey k, Functor f) => Trie f k a -> TrieZipper f k a
+zipper t = TrieZipper (t, Top)
+
+left, right, up :: (TrieKey k) => TrieZipper f k a -> Maybe (TrieZipper f k a)
+-- nowhere to go if we're at the leaf or in an empty tree
+left (getZipper -> (Empty, _)) = Nothing
+left (getZipper -> (Leaf _ _, _)) = Nothing
+left (getZipper -> (Branch bBit pref l r, c)) = Just $ TrieZipper (l, L bBit pref c r)
+-- nowhere to go if we're at the leaf or in an empty tree
+right (getZipper -> (Empty, _)) = Nothing
+right (getZipper -> (Leaf _ _, _)) = Nothing
+right (getZipper -> (Branch bBit pref l r, c)) = Just $ TrieZipper (r, R bBit pref l c)
+-- we're already at the top
+up (getZipper -> (_, Top)) = Nothing
+up (getZipper -> (t, L bBit pref c r)) = Just $ TrieZipper (Branch bBit pref t r, c)
+up (getZipper -> (t, R bBit pref l c)) = Just $ TrieZipper (Branch bBit pref l t, c)
+
+-- TODO not sure if this is what we want, need more research
+
+-- | Supplied function '(k -> a -> b)' takes key 'k' and value 'a' and returns 'b'
+-- upmostWith :: (TrieKey k, Functor f, Applicative f) => (k -> a -> b) -> TrieZipper f k a -> f (TrieZipper f k a)
+-- upmostWith _ z@(getZipper -> (_, Top)) = pure z
+-- upmostWith _ z@(getZipper -> (Empty, _)) = pure z
+-- upmostWith f z@(getZipper -> (Leaf key fa, c)) = undefined
+-- upmostWith f z@(getZipper -> (Branch bBit pref l r, c)) = undefined
+
+-- upmost :: (TrieKey k) => TrieZipper Identity k a -> TrieZipper Identity k a
+-- upmost = runIdentity . upmostWith const
+
+lookupZipper :: (TrieKey k, Functor f) => k -> Trie f k a -> Maybe (TrieZipper f k a)
+lookupZipper _ Empty = Nothing
+lookupZipper key1 t = go key1 $ TrieZipper (t, Top)
+  where
+    go :: (TrieKey k, Functor f) => k -> TrieZipper f k a -> Maybe (TrieZipper f k a)
+    go _ (getZipper -> (Empty, _)) = Nothing
+    go key z@(getZipper -> (Leaf k _, _)) = if key == k then Just z else Nothing
+    go key z@(getZipper -> (Branch bBit pref _ _, _))
+      | not (matchPrefix key pref bBit) = Nothing
+      | zeroBit key bBit = left z >>= go key
+      | otherwise = right z >>= go key
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index 2127795..8a819d8 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -60,6 +60,8 @@ library
     exposed-modules:
         Data.MerkleTrie
         Data.Trie
+        Data.Trie.Internal
+        Data.Trie.Zipper
 
     -- Modules included in this library but not exported.
     -- other-modules:

From 02899aef838a51d61e1095b9673880b783d551ae Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 16 Dec 2024 16:53:05 +0200
Subject: [PATCH 21/31] feat: zipper implementation [wip]

---
 merkle-patricia-trie/lib/Data/Trie.hs         | 31 +++++------
 .../lib/Data/Trie/Internal.hs                 | 48 ++++++++++------
 merkle-patricia-trie/lib/Data/Trie/Zipper.hs  | 55 +++++++++++--------
 merkle-patricia-trie/lib/Data/Utils.hs        | 27 +++++++++
 .../merkle-patricia-trie.cabal                |  4 +-
 5 files changed, 108 insertions(+), 57 deletions(-)
 create mode 100644 merkle-patricia-trie/lib/Data/Utils.hs

diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index cc638b5..a6a20b2 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,22 +1,21 @@
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.Trie (Trie (..), TrieKey, lookup, insert) where
+module Data.Trie (LittleTrie, BigTrie, TrieKey, lookup, insert, empty) where
 
-import Data.Bits
-import Data.Kind (Type)
 import Data.Trie.Internal
+import Data.Utils
 import Prelude hiding (lookup)
 
-lookup :: (TrieKey k, Functor f) => k -> Trie f k a -> Maybe (f a)
+-- TODO do we want to unwrap value here or return it wrapped? Maybe 2 different functions for these cases?
+lookup :: (TrieKey k, Wrapper w) => k -> Trie w k a -> Maybe (w a)
 lookup _ Empty = Nothing
 lookup k (Leaf key a) = if k == key then Just a else Nothing
-lookup k (Branch key prefix t1 t2)
-  | not (matchPrefix k prefix key) = Nothing
-  | zeroBit k key = lookup k t1
+lookup k (Branch bBit prefix t1 t2)
+  | not (matchPrefix k prefix bBit) = Nothing
+  | zeroBit k bBit = lookup k t1
   | otherwise = lookup k t2
 
--- maybe conflict func must be a Monoid?
-insertWith :: (TrieKey k, Functor f, Applicative f) => (v -> v -> v) -> k -> v -> Trie f k v -> Trie f k v
+insertWith :: (TrieKey k, Wrapper w) => (v -> v -> v) -> k -> v -> Trie w k v -> Trie w k v
 insertWith resolve k v t = go t
   where
     go Empty = Leaf k $ pure v
@@ -24,16 +23,16 @@ insertWith resolve k v t = go t
       if k == key
         then Leaf k $ resolve v <$> val
         else join k (Leaf k $ pure v) key t
-    go (Branch key prefix t1 t2) =
-      if matchPrefix k prefix key
+    go (Branch bBit prefix t1 t2) =
+      if matchPrefix k prefix bBit
         then
-          if zeroBit k key
-            then Branch key prefix (go t1) t2
-            else Branch key prefix t1 (go t2)
+          if zeroBit k bBit
+            then Branch bBit prefix (go t1) t2
+            else Branch bBit prefix t1 (go t2)
         else join k (Leaf k $ pure v) prefix t
 
-insert :: (TrieKey k, Functor f, Applicative f) => k -> v -> Trie f k v -> Trie f k v
+insert :: (TrieKey k, Wrapper w) => k -> v -> Trie w k v -> Trie w k v
 insert = insertWith const
 
-empty :: (TrieKey k, Functor f) => Trie f k v
+empty :: (TrieKey k, Wrapper w) => Trie w k v
 empty = Empty
diff --git a/merkle-patricia-trie/lib/Data/Trie/Internal.hs b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
index c1e6cd7..a2a343c 100644
--- a/merkle-patricia-trie/lib/Data/Trie/Internal.hs
+++ b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
@@ -1,25 +1,47 @@
+{-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module Data.Trie.Internal where
 
 import Data.Bits
 import Data.Kind (Type)
+import Data.WideWord (Word256)
+import Data.Word (Word8)
 import Prelude hiding (lookup)
 
-data Trie (f :: Type -> Type) k a
+data Trie' (w :: Type -> Type) k h a
   = Empty
-  | Leaf k (f a)
-  | Branch k k (Trie f k a) (Trie f k a)
+  | Leaf k (w a)
+  | Branch k k (Trie' w k h a) (Trie' w k h a)
+  deriving (Functor)
+
+-- | This will allow us to enforce the height of the trie depending on it's key
+type family TrieHeight key where
+  TrieHeight (LEK Word256) = Word8
+  TrieHeight (BEK Word256) = Word8
+
+-- | This will ensure, that when someone tries to create a Trie, the height will be automagically supplied
+type Trie w k a = Trie' w k (TrieHeight k) a
+
+type LittleTrie w k a = Trie w (LEK k) a
+
+type BigTrie w k a = Trie w (BEK k) a
 
 class (Bits k, Num k, Integral k) => TrieKey k where
   -- | Discards bits before/after(depending on endianess) branching bit
   mask :: k -> k -> k
 
-  -- | Leaves only the branching bit(so key 0b1001101 becomes 0b0001000, if 4th bit was branching)
+  -- | Leaves only the branching bit(so key 0b1001101 becomes 0b0001000, if 4th bit was a branching bit)
   branchingBit :: k -> k -> k
 
+  matchPrefix :: k -> k -> k -> Bool
+  matchPrefix bBit prefix key = mask bBit key == prefix
+
+  zeroBit :: k -> k -> Bool
+  zeroBit a b = (a .&. b) == 0
+
 newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
 
 newtype BEK a = BigEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
@@ -37,16 +59,10 @@ instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
   -- since it needs more args than just a and b(it also needs prefixes)
   branchingBit _ _ = undefined
 
-matchPrefix :: (TrieKey k) => k -> k -> k -> Bool
-matchPrefix k prefix key = mask k key == prefix
-
-zeroBit :: (TrieKey k) => k -> k -> Bool
-zeroBit a b = (a .&. b) == 0
-
-join :: (TrieKey k) => k -> Trie f k a -> k -> Trie f k a -> Trie f k a
+join :: (TrieKey k) => k -> Trie w k a -> k -> Trie w k a -> Trie w k a
 join p1 t1 p2 t2 =
-  if zeroBit p1 commonPrefix
-    then Branch (mask p1 commonPrefix) commonPrefix t1 t2
-    else Branch (mask p1 commonPrefix) commonPrefix t2 t1
+  if zeroBit p1 bBit
+    then Branch bBit (mask p1 bBit) t1 t2
+    else Branch bBit (mask p1 bBit) t2 t1
   where
-    commonPrefix = branchingBit p1 p2
+    bBit = branchingBit p1 p2
diff --git a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
index 0edd40e..31a45f6 100644
--- a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
+++ b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
@@ -1,23 +1,26 @@
+{-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE ViewPatterns #-}
 
-module Data.Trie.Zipper (TrieZipper, left, right, up, lookupZipper, zipper) where
+module Data.Trie.Zipper (TrieZipper, left, right, up, focus, zipper) where
 
--- import Control.Monad.Identity (Identity (..))
 import Data.Kind (Type)
+import Data.Maybe (fromMaybe)
 import Data.Trie.Internal
+import Data.Utils
 
-data Cxt (f :: Type -> Type) k a
+data Cxt (w :: Type -> Type) k a
   = Top
-  | L k k (Cxt f k a) (Trie f k a)
-  | R k k (Trie f k a) (Cxt f k a)
+  | L k k (Cxt w k a) (Trie w k a)
+  | R k k (Trie w k a) (Cxt w k a)
+  deriving (Functor)
 
-newtype TrieZipper (f :: Type -> Type) k a = TrieZipper {getZipper :: (Trie f k a, Cxt f k a)}
+newtype TrieZipper (w :: Type -> Type) k a = TrieZipper {getZipper :: (Trie w k a, Cxt w k a)} deriving (Functor)
 
-zipper :: (TrieKey k, Functor f) => Trie f k a -> TrieZipper f k a
+zipper :: (TrieKey k, Wrapper w) => Trie w k a -> TrieZipper w k a
 zipper t = TrieZipper (t, Top)
 
-left, right, up :: (TrieKey k) => TrieZipper f k a -> Maybe (TrieZipper f k a)
+left, right, up :: (TrieKey k) => TrieZipper w k a -> Maybe (TrieZipper w k a)
 -- nowhere to go if we're at the leaf or in an empty tree
 left (getZipper -> (Empty, _)) = Nothing
 left (getZipper -> (Leaf _ _, _)) = Nothing
@@ -31,23 +34,27 @@ up (getZipper -> (_, Top)) = Nothing
 up (getZipper -> (t, L bBit pref c r)) = Just $ TrieZipper (Branch bBit pref t r, c)
 up (getZipper -> (t, R bBit pref l c)) = Just $ TrieZipper (Branch bBit pref l t, c)
 
--- TODO not sure if this is what we want, need more research
-
--- | Supplied function '(k -> a -> b)' takes key 'k' and value 'a' and returns 'b'
--- upmostWith :: (TrieKey k, Functor f, Applicative f) => (k -> a -> b) -> TrieZipper f k a -> f (TrieZipper f k a)
--- upmostWith _ z@(getZipper -> (_, Top)) = pure z
--- upmostWith _ z@(getZipper -> (Empty, _)) = pure z
--- upmostWith f z@(getZipper -> (Leaf key fa, c)) = undefined
--- upmostWith f z@(getZipper -> (Branch bBit pref l r, c)) = undefined
-
--- upmost :: (TrieKey k) => TrieZipper Identity k a -> TrieZipper Identity k a
--- upmost = runIdentity . upmostWith const
-
-lookupZipper :: (TrieKey k, Functor f) => k -> Trie f k a -> Maybe (TrieZipper f k a)
-lookupZipper _ Empty = Nothing
-lookupZipper key1 t = go key1 $ TrieZipper (t, Top)
+modify :: (Wrapper w) => (Trie w k a -> Trie w k a) -> TrieZipper w k a -> TrieZipper w k a
+modify f (getZipper -> (t, c)) = TrieZipper (f t, c)
+
+-- | Move focus to the top of the trie applying `(k -> a -> a)` along the way
+-- Supplied function '(k -> a -> a)' takes key 'k' and value 'a' and returns 'a'
+unfocusWith :: (TrieKey k, Wrapper w) => (k -> a -> b) -> TrieZipper w k a -> TrieZipper w k b
+unfocusWith f z@(getZipper -> (_, Top)) = fmap f z
+unfocusWith f z@(getZipper -> (Empty, _)) = fmap f z
+unfocusWith f z@(getZipper -> (Leaf key fa, c)) = unfocusWith f . fromMaybe z . up $ TrieZipper (Leaf key $ fmap (f key) fa, fmap f c)
+unfocusWith f z@(getZipper -> (Branch bBit pref l r, c)) = unfocusWith f . fromMaybe z $ up z
+
+-- | Move focus to the top of the trie
+unfocus :: (TrieKey k, Wrapper w) => TrieZipper w k a -> TrieZipper w k a
+unfocus = upmostWith (const id)
+
+-- | Create a `TrieZipper` focused on given key `k` in a trie `Trie w k a`
+focus :: (TrieKey k, Wrapper w) => k -> Trie w k a -> Maybe (TrieZipper w k a)
+focus _ Empty = Nothing
+focus key1 t = go key1 $ TrieZipper (t, Top)
   where
-    go :: (TrieKey k, Functor f) => k -> TrieZipper f k a -> Maybe (TrieZipper f k a)
+    go :: (TrieKey k, Wrapper w) => k -> TrieZipper w k a -> Maybe (TrieZipper w k a)
     go _ (getZipper -> (Empty, _)) = Nothing
     go key z@(getZipper -> (Leaf k _, _)) = if key == k then Just z else Nothing
     go key z@(getZipper -> (Branch bBit pref _ _, _))
diff --git a/merkle-patricia-trie/lib/Data/Utils.hs b/merkle-patricia-trie/lib/Data/Utils.hs
new file mode 100644
index 0000000..27803d1
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Utils.hs
@@ -0,0 +1,27 @@
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE TypeOperators #-}
+
+module Data.Utils where
+
+import Data.Functor.Const
+import Data.Functor.Identity
+
+-- | Type of natural transformation between 'f' and 'g'
+type (~>) f g = forall x. f x -> g x
+
+-- | Typeclass for extractible wrappers 'w'
+class (Functor w, Applicative w) => Wrapper w where
+  unwrap :: w a -> a
+
+  wrap :: a -> w a
+  wrap = pure
+
+instance Wrapper Identity where
+  unwrap = runIdentity
+
+-- Not sure if this can even be used
+class (Wrapper w1, Wrapper w2) => Transform w1 w2 where
+  transform :: w1 a -> w2 b
+
+-- class (Functor f, Functor g) => NaturalTransformation f g where
+--   alpha :: f a -> g a
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index 8a819d8..54b2e0d 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -62,6 +62,7 @@ library
         Data.Trie
         Data.Trie.Internal
         Data.Trie.Zipper
+        Data.Utils
 
     -- Modules included in this library but not exported.
     -- other-modules:
@@ -74,7 +75,8 @@ library
     build-depends:
         base ^>=4.18.2.1,
         cryptonite ==0.30,
-        bytestring
+        bytestring,
+        wide-word
         -- for future more optimized implementation consider using RLP from ethereum whitepaper
         -- relapse ==1.0.0.1
 

From ac927a7460816fdc6b4c771e08366fe1ce18359b Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Tue, 17 Dec 2024 16:26:54 +0200
Subject: [PATCH 22/31] feat: refactoring of trie&zipper implementations

---
 .../lib/Data/Trie/BigEndian.hs                | 26 +++++++++++
 .../lib/Data/Trie/Internal.hs                 | 44 +++++--------------
 .../lib/Data/Trie/LittleEndian.hs             | 25 +++++++++++
 merkle-patricia-trie/lib/Data/Utils.hs        | 26 -----------
 .../merkle-patricia-trie.cabal                |  2 +
 5 files changed, 63 insertions(+), 60 deletions(-)
 create mode 100644 merkle-patricia-trie/lib/Data/Trie/BigEndian.hs
 create mode 100644 merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs

diff --git a/merkle-patricia-trie/lib/Data/Trie/BigEndian.hs b/merkle-patricia-trie/lib/Data/Trie/BigEndian.hs
new file mode 100644
index 0000000..c2932c3
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Trie/BigEndian.hs
@@ -0,0 +1,26 @@
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Data.Trie.BigEndian where
+
+import Data.Bits
+import Data.Kind (Type)
+import Data.Trie.Internal
+import Data.WideWord (Word256)
+import Data.Word (Word8)
+import Prelude hiding (lookup)
+
+newtype BEK a = BigEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
+
+newtype Trie k a = Trie {unTrie :: Trie' (BEK k) (TrieHeight k) a} deriving (Show, Eq, Functor)
+
+instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
+  mask (BigEndian a) (BigEndian b) = BigEndian $ (a .|. (b - 1)) .&. complement b
+
+  -- TODO figure out how to properly implement branchingBit for BigEndian
+  -- since it needs more args than just a and b(it also needs prefixes)
+  branchingBit _ _ = undefined
diff --git a/merkle-patricia-trie/lib/Data/Trie/Internal.hs b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
index a2a343c..ad7f429 100644
--- a/merkle-patricia-trie/lib/Data/Trie/Internal.hs
+++ b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
 
@@ -11,25 +13,16 @@ import Data.WideWord (Word256)
 import Data.Word (Word8)
 import Prelude hiding (lookup)
 
-data Trie' (w :: Type -> Type) k h a
+data Trie' k h a
   = Empty
-  | Leaf k (w a)
-  | Branch k k (Trie' w k h a) (Trie' w k h a)
-  deriving (Functor)
-
--- | This will allow us to enforce the height of the trie depending on it's key
-type family TrieHeight key where
-  TrieHeight (LEK Word256) = Word8
-  TrieHeight (BEK Word256) = Word8
-
--- | This will ensure, that when someone tries to create a Trie, the height will be automagically supplied
-type Trie w k a = Trie' w k (TrieHeight k) a
-
-type LittleTrie w k a = Trie w (LEK k) a
-
-type BigTrie w k a = Trie w (BEK k) a
+  | Leaf {key :: k, val :: a}
+  | Branch {bBit :: h, pref :: k, left :: Trie' k h a, right :: Trie' k h a}
+  deriving (Show, Eq, Functor)
 
 class (Bits k, Num k, Integral k) => TrieKey k where
+  -- | This will allow us to enforce the height of the trie depending on it's key
+  type TrieHeight k :: Type
+
   -- | Discards bits before/after(depending on endianess) branching bit
   mask :: k -> k -> k
 
@@ -42,24 +35,7 @@ class (Bits k, Num k, Integral k) => TrieKey k where
   zeroBit :: k -> k -> Bool
   zeroBit a b = (a .&. b) == 0
 
-newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
-
-newtype BEK a = BigEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
-
-instance (Bits k, Num k, Integral k) => TrieKey (LEK k) where
-  mask (LittleEndian a) (LittleEndian b) = LittleEndian $ a .&. (b - 1)
-  branchingBit (LittleEndian a) (LittleEndian b) = LittleEndian . lowestBit $ a `xor` b
-    where
-      lowestBit x = x .&. complement x
-
-instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
-  mask (BigEndian a) (BigEndian b) = BigEndian $ (a .|. (b - 1)) .&. complement b
-
-  -- TODO figure out how to properly implement branchingBit for BigEndian
-  -- since it needs more args than just a and b(it also needs prefixes)
-  branchingBit _ _ = undefined
-
-join :: (TrieKey k) => k -> Trie w k a -> k -> Trie w k a -> Trie w k a
+join :: (TrieKey k) => k -> Trie' k h a -> k -> Trie' k h a -> Trie' k h a
 join p1 t1 p2 t2 =
   if zeroBit p1 bBit
     then Branch bBit (mask p1 bBit) t1 t2
diff --git a/merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs b/merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs
new file mode 100644
index 0000000..37ae140
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs
@@ -0,0 +1,25 @@
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Data.Trie.LittleEndian where
+
+import Data.Bits
+import Data.Kind (Type)
+import Data.Trie.Internal
+import Data.WideWord (Word256)
+import Data.Word (Word8)
+import Prelude hiding (lookup)
+
+newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
+
+newtype Trie k a = Trie {unTrie :: Trie' (LEK k) (TrieHeight k) a} deriving (Show, Eq, Functor)
+
+instance (Bits k, Num k, Integral k) => TrieKey (LEK k) where
+  mask (LittleEndian a) (LittleEndian b) = LittleEndian $ a .&. (b - 1)
+  branchingBit (LittleEndian a) (LittleEndian b) = LittleEndian . lowestBit $ a `xor` b
+    where
+      lowestBit x = x .&. complement x
diff --git a/merkle-patricia-trie/lib/Data/Utils.hs b/merkle-patricia-trie/lib/Data/Utils.hs
index 27803d1..19e622d 100644
--- a/merkle-patricia-trie/lib/Data/Utils.hs
+++ b/merkle-patricia-trie/lib/Data/Utils.hs
@@ -1,27 +1 @@
-{-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE TypeOperators #-}
-
 module Data.Utils where
-
-import Data.Functor.Const
-import Data.Functor.Identity
-
--- | Type of natural transformation between 'f' and 'g'
-type (~>) f g = forall x. f x -> g x
-
--- | Typeclass for extractible wrappers 'w'
-class (Functor w, Applicative w) => Wrapper w where
-  unwrap :: w a -> a
-
-  wrap :: a -> w a
-  wrap = pure
-
-instance Wrapper Identity where
-  unwrap = runIdentity
-
--- Not sure if this can even be used
-class (Wrapper w1, Wrapper w2) => Transform w1 w2 where
-  transform :: w1 a -> w2 b
-
--- class (Functor f, Functor g) => NaturalTransformation f g where
---   alpha :: f a -> g a
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index 54b2e0d..f28e42f 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -61,6 +61,8 @@ library
         Data.MerkleTrie
         Data.Trie
         Data.Trie.Internal
+        Data.Trie.LittleEndian
+        Data.Trie.BigEndian
         Data.Trie.Zipper
         Data.Utils
 

From 79f1e82e455e4963e6d32dcac5562421d4db4caa Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Thu, 19 Dec 2024 16:46:17 +0200
Subject: [PATCH 23/31] feat: refactoring of trie&zipper implementation[part2]

---
 merkle-patricia-trie/lib/Data/Trie.hs         | 43 ++++++++--------
 .../lib/Data/Trie/BigEndian.hs                | 26 ----------
 .../lib/Data/Trie/Internal.hs                 | 50 ++++++++++++++-----
 .../lib/Data/Trie/LittleEndian.hs             | 25 ----------
 merkle-patricia-trie/lib/Data/Trie/Zipper.hs  | 49 ++++++++++--------
 .../merkle-patricia-trie.cabal                |  2 -
 6 files changed, 87 insertions(+), 108 deletions(-)
 delete mode 100644 merkle-patricia-trie/lib/Data/Trie/BigEndian.hs
 delete mode 100644 merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs

diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index a6a20b2..69ee3b7 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,38 +1,39 @@
+{-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.Trie (LittleTrie, BigTrie, TrieKey, lookup, insert, empty) where
+module Data.Trie (Trie, TrieKey, lookup, insertWith, insert, empty) where
 
 import Data.Trie.Internal
-import Data.Utils
 import Prelude hiding (lookup)
 
 -- TODO do we want to unwrap value here or return it wrapped? Maybe 2 different functions for these cases?
-lookup :: (TrieKey k, Wrapper w) => k -> Trie w k a -> Maybe (w a)
+lookup :: (TrieKey k) => k -> Trie k a -> Maybe a
 lookup _ Empty = Nothing
-lookup k (Leaf key a) = if k == key then Just a else Nothing
-lookup k (Branch bBit prefix t1 t2)
-  | not (matchPrefix k prefix bBit) = Nothing
-  | zeroBit k bBit = lookup k t1
-  | otherwise = lookup k t2
+lookup k Leaf {..} = if k == key then Just value else Nothing
+lookup k Branch {..}
+  | not (matchPrefix k prefix (heightToBBit branchingBit)) = Nothing
+  | zeroBit k (heightToBBit branchingBit) = lookup k left
+  | otherwise = lookup k right
 
-insertWith :: (TrieKey k, Wrapper w) => (v -> v -> v) -> k -> v -> Trie w k v -> Trie w k v
+insertWith :: forall k a. (TrieKey k) => (a -> a -> a) -> k -> a -> Trie k a -> Trie k a
 insertWith resolve k v t = go t
   where
-    go Empty = Leaf k $ pure v
-    go (Leaf key val) =
+    go :: Trie k a -> Trie k a
+    go Empty = Leaf k v
+    go Leaf {..} =
       if k == key
-        then Leaf k $ resolve v <$> val
-        else join k (Leaf k $ pure v) key t
-    go (Branch bBit prefix t1 t2) =
-      if matchPrefix k prefix bBit
+        then Leaf k $ resolve v value
+        else join k (Leaf k v) key t
+    go Branch {..} =
+      if matchPrefix k prefix (heightToBBit branchingBit)
         then
-          if zeroBit k bBit
-            then Branch bBit prefix (go t1) t2
-            else Branch bBit prefix t1 (go t2)
-        else join k (Leaf k $ pure v) prefix t
+          if zeroBit k (heightToBBit branchingBit)
+            then Branch branchingBit prefix (go left) right
+            else Branch branchingBit prefix left (go right)
+        else join k (Leaf k v) prefix t
 
-insert :: (TrieKey k, Wrapper w) => k -> v -> Trie w k v -> Trie w k v
+insert :: (TrieKey k) => k -> v -> Trie k v -> Trie k v
 insert = insertWith const
 
-empty :: (TrieKey k, Wrapper w) => Trie w k v
+empty :: (TrieKey k) => Trie k v
 empty = Empty
diff --git a/merkle-patricia-trie/lib/Data/Trie/BigEndian.hs b/merkle-patricia-trie/lib/Data/Trie/BigEndian.hs
deleted file mode 100644
index c2932c3..0000000
--- a/merkle-patricia-trie/lib/Data/Trie/BigEndian.hs
+++ /dev/null
@@ -1,26 +0,0 @@
-{-# LANGUAGE DeriveFunctor #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE OverloadedRecordDot #-}
-{-# LANGUAGE RecordWildCards #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeFamilies #-}
-
-module Data.Trie.BigEndian where
-
-import Data.Bits
-import Data.Kind (Type)
-import Data.Trie.Internal
-import Data.WideWord (Word256)
-import Data.Word (Word8)
-import Prelude hiding (lookup)
-
-newtype BEK a = BigEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
-
-newtype Trie k a = Trie {unTrie :: Trie' (BEK k) (TrieHeight k) a} deriving (Show, Eq, Functor)
-
-instance (Bits k, Num k, Integral k) => TrieKey (BEK k) where
-  mask (BigEndian a) (BigEndian b) = BigEndian $ (a .|. (b - 1)) .&. complement b
-
-  -- TODO figure out how to properly implement branchingBit for BigEndian
-  -- since it needs more args than just a and b(it also needs prefixes)
-  branchingBit _ _ = undefined
diff --git a/merkle-patricia-trie/lib/Data/Trie/Internal.hs b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
index ad7f429..38bc72e 100644
--- a/merkle-patricia-trie/lib/Data/Trie/Internal.hs
+++ b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE RecordWildCards #-}
@@ -13,32 +14,57 @@ import Data.WideWord (Word256)
 import Data.Word (Word8)
 import Prelude hiding (lookup)
 
+-- 0000010
+
 data Trie' k h a
   = Empty
-  | Leaf {key :: k, val :: a}
-  | Branch {bBit :: h, pref :: k, left :: Trie' k h a, right :: Trie' k h a}
+  | Leaf {key :: k, value :: a}
+  | -- | Branch node stores branching bit `bBit` and longest common prefix `pref`
+    Branch {branchingBit :: h, prefix :: k, left :: Trie' k h a, right :: Trie' k h a}
   deriving (Show, Eq, Functor)
 
-class (Bits k, Num k, Integral k) => TrieKey k where
+-- | Constraint for Trie key and height. Includes default implementation for little endian Tries
+-- TODO replace default with big endian Trie
+class (Bits k, Num k, Integral k, Integral (TrieHeight k)) => TrieKey k where
   -- | This will allow us to enforce the height of the trie depending on it's key
   type TrieHeight k :: Type
 
-  -- | Discards bits before/after(depending on endianess) branching bit
+  -- | Discards bits not in the mask
   mask :: k -> k -> k
+  mask k m = k .&. (m - 1)
 
-  -- | Leaves only the branching bit(so key 0b1001101 becomes 0b0001000, if 4th bit was a branching bit)
-  branchingBit :: k -> k -> k
+  -- | Finds the first bit on which prefixes disagree.
+  commonBranchingBit :: k -> k -> k
+  commonBranchingBit p1 p2 = lowestBit (p1 `xor` p2)
+    where
+      lowestBit :: k -> k
+      lowestBit x = x .&. complement x
 
+  -- | Masks key using supplied branching bit and compares to prefix
   matchPrefix :: k -> k -> k -> Bool
-  matchPrefix bBit prefix key = mask bBit key == prefix
+  matchPrefix key prefix m = mask key m == prefix
+
+  -- | Converts height 'TrieHeight k' into the branching bit 'k'
+  heightToBBit :: TrieHeight k -> k
+  heightToBBit h = 2 ^ h
+
+  -- | Converts branching bit 'k' into height 'TrieHeight k'
+  bBitToHeight :: k -> TrieHeight k
+  bBitToHeight = floor . logBase 2.0 . fromIntegral
 
+  -- | Tests whether the desired bit is zero
   zeroBit :: k -> k -> Bool
-  zeroBit a b = (a .&. b) == 0
+  zeroBit k m = (k .&. m) == 0
+
+type Trie k a = Trie' k (TrieHeight k) a
+
+branch :: (TrieKey k) => k -> k -> Trie k a -> Trie k a -> Trie k a
+branch bBit = Branch (bBitToHeight bBit)
 
-join :: (TrieKey k) => k -> Trie' k h a -> k -> Trie' k h a -> Trie' k h a
+join :: (TrieKey k) => k -> Trie k a -> k -> Trie k a -> Trie k a
 join p1 t1 p2 t2 =
   if zeroBit p1 bBit
-    then Branch bBit (mask p1 bBit) t1 t2
-    else Branch bBit (mask p1 bBit) t2 t1
+    then branch bBit (mask p1 bBit) t1 t2
+    else branch bBit (mask p1 bBit) t2 t1
   where
-    bBit = branchingBit p1 p2
+    bBit = commonBranchingBit p1 p2
diff --git a/merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs b/merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs
deleted file mode 100644
index 37ae140..0000000
--- a/merkle-patricia-trie/lib/Data/Trie/LittleEndian.hs
+++ /dev/null
@@ -1,25 +0,0 @@
-{-# LANGUAGE DeriveFunctor #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE OverloadedRecordDot #-}
-{-# LANGUAGE RecordWildCards #-}
-{-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeFamilies #-}
-
-module Data.Trie.LittleEndian where
-
-import Data.Bits
-import Data.Kind (Type)
-import Data.Trie.Internal
-import Data.WideWord (Word256)
-import Data.Word (Word8)
-import Prelude hiding (lookup)
-
-newtype LEK a = LittleEndian a deriving (Enum, Eq, Ord, Real, Bits, Num, Integral)
-
-newtype Trie k a = Trie {unTrie :: Trie' (LEK k) (TrieHeight k) a} deriving (Show, Eq, Functor)
-
-instance (Bits k, Num k, Integral k) => TrieKey (LEK k) where
-  mask (LittleEndian a) (LittleEndian b) = LittleEndian $ a .&. (b - 1)
-  branchingBit (LittleEndian a) (LittleEndian b) = LittleEndian . lowestBit $ a `xor` b
-    where
-      lowestBit x = x .&. complement x
diff --git a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
index 31a45f6..ab11393 100644
--- a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
+++ b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ViewPatterns #-}
 
 module Data.Trie.Zipper (TrieZipper, left, right, up, focus, zipper) where
@@ -7,57 +8,61 @@ module Data.Trie.Zipper (TrieZipper, left, right, up, focus, zipper) where
 import Data.Kind (Type)
 import Data.Maybe (fromMaybe)
 import Data.Trie.Internal
-import Data.Utils
 
-data Cxt (w :: Type -> Type) k a
+data Cxt (f :: Type -> Type) k h a
   = Top
-  | L k k (Cxt w k a) (Trie w k a)
-  | R k k (Trie w k a) (Cxt w k a)
+  | L k h (f a) (Cxt f k h a) (Trie k a)
+  | R k h (f a) (Trie k a) (Cxt f k h a)
   deriving (Functor)
 
-newtype TrieZipper (w :: Type -> Type) k a = TrieZipper {getZipper :: (Trie w k a, Cxt w k a)} deriving (Functor)
+newtype TrieZipper' (f :: Type -> Type) k h a = TrieZipper {getZipper :: (Trie k a, Cxt f k h a)} deriving (Functor)
 
-zipper :: (TrieKey k, Wrapper w) => Trie w k a -> TrieZipper w k a
+type TrieZipper (f :: Type -> Type) k a = TrieZipper' f k (TrieHeight k) a
+
+zipper :: (TrieKey k, Functor f) => Trie k a -> TrieZipper f k a
 zipper t = TrieZipper (t, Top)
 
-left, right, up :: (TrieKey k) => TrieZipper w k a -> Maybe (TrieZipper w k a)
+-- TODO do we want to return Maybe zipper or just return the same one if the move is errorneous?
+-- Maybe 2 family of functions for this?
+left, right, up :: (TrieKey k, Applicative f) => TrieZipper f k a -> Maybe (TrieZipper f k a)
 -- nowhere to go if we're at the leaf or in an empty tree
 left (getZipper -> (Empty, _)) = Nothing
-left (getZipper -> (Leaf _ _, _)) = Nothing
-left (getZipper -> (Branch bBit pref l r, c)) = Just $ TrieZipper (l, L bBit pref c r)
+left (getZipper -> (Leaf {}, _)) = Nothing
+-- TODO how to calculate values here? For leaf it's easy, as for branch? It must depend on the leafes nearby, I guess
+left (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (lelft, L branchingBit prefix undefined c right)
 -- nowhere to go if we're at the leaf or in an empty tree
 right (getZipper -> (Empty, _)) = Nothing
 right (getZipper -> (Leaf _ _, _)) = Nothing
-right (getZipper -> (Branch bBit pref l r, c)) = Just $ TrieZipper (r, R bBit pref l c)
+right (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (right, R branchingBit prefex undefined left c)
 -- we're already at the top
 up (getZipper -> (_, Top)) = Nothing
-up (getZipper -> (t, L bBit pref c r)) = Just $ TrieZipper (Branch bBit pref t r, c)
-up (getZipper -> (t, R bBit pref l c)) = Just $ TrieZipper (Branch bBit pref l t, c)
+up (getZipper -> (t, L branchingBit prefix v c r)) = Just $ TrieZipper (Branch branchingBit prefix t r, c)
+up (getZipper -> (t, R branchingBit prefix v l c)) = Just $ TrieZipper (Branch branchingBit prefix l t, c)
 
-modify :: (Wrapper w) => (Trie w k a -> Trie w k a) -> TrieZipper w k a -> TrieZipper w k a
+modify :: (Applicative f) => (Trie k a -> Trie k a) -> TrieZipper f k a -> TrieZipper f k a
 modify f (getZipper -> (t, c)) = TrieZipper (f t, c)
 
 -- | Move focus to the top of the trie applying `(k -> a -> a)` along the way
 -- Supplied function '(k -> a -> a)' takes key 'k' and value 'a' and returns 'a'
-unfocusWith :: (TrieKey k, Wrapper w) => (k -> a -> b) -> TrieZipper w k a -> TrieZipper w k b
+unfocusWith :: (TrieKey k, Applicative f) => (k -> a -> b) -> TrieZipper f k a -> TrieZipper f k b
 unfocusWith f z@(getZipper -> (_, Top)) = fmap f z
 unfocusWith f z@(getZipper -> (Empty, _)) = fmap f z
 unfocusWith f z@(getZipper -> (Leaf key fa, c)) = unfocusWith f . fromMaybe z . up $ TrieZipper (Leaf key $ fmap (f key) fa, fmap f c)
 unfocusWith f z@(getZipper -> (Branch bBit pref l r, c)) = unfocusWith f . fromMaybe z $ up z
 
 -- | Move focus to the top of the trie
-unfocus :: (TrieKey k, Wrapper w) => TrieZipper w k a -> TrieZipper w k a
+unfocus :: (TrieKey k, Applicative f) => TrieZipper f k a -> TrieZipper f k a
 unfocus = upmostWith (const id)
 
 -- | Create a `TrieZipper` focused on given key `k` in a trie `Trie w k a`
-focus :: (TrieKey k, Wrapper w) => k -> Trie w k a -> Maybe (TrieZipper w k a)
+focus :: (TrieKey k, Applicative f) => k -> Trie k a -> Maybe (TrieZipper f k a)
 focus _ Empty = Nothing
 focus key1 t = go key1 $ TrieZipper (t, Top)
   where
-    go :: (TrieKey k, Wrapper w) => k -> TrieZipper w k a -> Maybe (TrieZipper w k a)
+    go :: (TrieKey k, Functor f) => k -> TrieZipper f k a -> Maybe (TrieZipper f k a)
     go _ (getZipper -> (Empty, _)) = Nothing
-    go key z@(getZipper -> (Leaf k _, _)) = if key == k then Just z else Nothing
-    go key z@(getZipper -> (Branch bBit pref _ _, _))
-      | not (matchPrefix key pref bBit) = Nothing
-      | zeroBit key bBit = left z >>= go key
-      | otherwise = right z >>= go key
+    go k z@(getZipper -> (Leaf {..}, _)) = if key == k then Just z else Nothing
+    go k z@(getZipper -> (Branch {..}, _))
+      | not (matchPrefix k prefix (heightToBBit branchingBit)) = Nothing
+      | zeroBit k (heightToBBit branchingBit) = left z >>= go k
+      | otherwise = right z >>= go k
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index f28e42f..54b2e0d 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -61,8 +61,6 @@ library
         Data.MerkleTrie
         Data.Trie
         Data.Trie.Internal
-        Data.Trie.LittleEndian
-        Data.Trie.BigEndian
         Data.Trie.Zipper
         Data.Utils
 

From 0af45dd33214f6db11c01437f47cac17e5e5359c Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Tue, 24 Dec 2024 02:01:26 +0200
Subject: [PATCH 24/31] feat: merkle proof impl

---
 merkle-patricia-trie/lib/Data/MerkleTrie.hs   | 122 ++++--------------
 merkle-patricia-trie/lib/Data/Trie.hs         |  20 +--
 .../lib/Data/Trie/Internal.hs                 |  13 +-
 merkle-patricia-trie/lib/Data/Trie/Zipper.hs  | 119 +++++++++++------
 .../merkle-patricia-trie.cabal                |   9 +-
 5 files changed, 129 insertions(+), 154 deletions(-)

diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index b1deaf9..010c762 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -1,103 +1,35 @@
-{-# LANGUAGE DeriveGeneric #-}
-{-# LANGUAGE NamedFieldPuns #-}
-{-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
 
 module Data.MerkleTrie where
 
-import Control.Applicative ((<|>))
 import Crypto.Hash (Digest, SHA256, hash)
 import Data.ByteString.Char8 (pack)
-import Data.List (isPrefixOf)
-import GHC.Generics (Generic)
-
--- | Patricia Trie node
-data PatriciaTrie k v = Node
-  { prefix :: k,
-    value :: Maybe v,
-    children :: [PatriciaTrie k v]
-  }
-  deriving (Show, Generic)
-
--- | Merkle Patricia Trie
-data MerklePatriciaTrie k v = MerklePatriciaTrie
-  { rootHash :: Digest SHA256,
-    rootNode :: PatriciaTrie k v
-  }
-  deriving (Show, Generic)
-
--- | Compute the hash of a Patricia Trie node
-computeNodeHash :: (Show k, Show v) => PatriciaTrie k v -> Digest SHA256
-computeNodeHash Node {..} =
-  hash . pack $ show (prefix, value, map computeNodeHash children)
-
--- | Update the root hash of the Merkle Trie
-updateRootHash :: (Show k, Show v) => PatriciaTrie k v -> Digest SHA256
-updateRootHash = computeNodeHash
-
--- | Create an empty Merkle Patricia Trie
-emptyTrie :: forall k v. (Show k, Show v) => MerklePatriciaTrie [k] v
-emptyTrie =
-  MerklePatriciaTrie
-    { rootHash = computeNodeHash @[k] @v root,
-      rootNode = root
-    }
+import Data.Trie
+
+computeHash :: (Show a) => a -> Digest SHA256
+computeHash = hash . pack . show
+
+-- | Compute the hash of a Patricia Trie
+-- | For internal use only, since it's partially matched against the trie!
+-- TODO don't like the fact, that this is an incomplete pattern, although I can wrap this in Maybe
+-- need to think about this
+computeNodeHash :: (Show k, Show (TrieHeight k), Show a) => Trie k a -> Digest SHA256
+computeNodeHash Leaf {..} = computeHash (key, value)
+computeNodeHash Branch {..} = computeHash (computeNodeHash left, computeNodeHash right)
+
+prove :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => k -> Trie k a -> Maybe (Digest SHA256, [Digest SHA256])
+prove k t = do
+  z <- focus k t
+  startingHash <- computeHash <$> getKeyValue z -- compute hash for the leaf
+  let accum = (startingHash, []) -- initial accumulator
+      zUp = goUp z -- now we're at the bottom-most branch
+  pure . fst $ upmostWith go accum zUp
   where
-    root = Node {prefix = [], value = Nothing, children = []}
-
--- | Create a singleton Patricia Trie
-singletonTrie :: forall k v. (Show k, Show v, Eq k) => [k] -> v -> MerklePatriciaTrie [k] v
-singletonTrie k v =
-  MerklePatriciaTrie
-    { rootHash = computeNodeHash @[k] @v root,
-      rootNode = root
-    }
-  where
-    root = Node {prefix = k, value = Just v, children = []}
-
--- | Insert a key-value pair into a Patricia Trie node
-insertNode :: (Eq k, Ord k) => [k] -> v -> PatriciaTrie [k] v -> PatriciaTrie [k] v
-insertNode [] v node = node {value = Just v}
-insertNode key v Node {..}
-  | prefix `isPrefixOf` key =
-      let remaining = drop (length prefix) key
-       in if null remaining
-            then Node {value = Just v, ..}
-            else Node {children = insertChild remaining v children, ..}
-  | otherwise = mergeNodes key v prefix value children
-  where
-    insertChild [] val [] = [Node {prefix = [], value = Just val, children = []}]
-    insertChild r val (c : cs)
-      | commonPrefix r c.prefix /= [] =
-          let splitChild = insertNode (drop (length (commonPrefix r c.prefix)) r) val c
-           in splitChild : cs
-      | otherwise = c : insertChild r val cs
-    insertChild r val [] = [Node {prefix = r, value = Just val, children = []}]
-    mergeNodes newKey newValue oldKey oldValue oldChildren =
-      let common = commonPrefix newKey oldKey
-          splitOld = Node {prefix = drop (length common) oldKey, value = oldValue, children = oldChildren}
-          splitNew = Node {prefix = drop (length common) newKey, value = Just newValue, children = []}
-       in Node {prefix = common, value = Nothing, children = [splitOld, splitNew]}
-    commonPrefix :: (Eq a) => [a] -> [a] -> [a]
-    commonPrefix [] _ = []
-    commonPrefix _ [] = []
-    commonPrefix (x : xs) (y : ys)
-      | x == y = x : commonPrefix xs ys
-      | otherwise = []
-
--- | Insert into the Merkle Patricia Trie
-insertTrie :: (Show k, Ord k, Eq k, Show v) => [k] -> v -> MerklePatriciaTrie [k] v -> MerklePatriciaTrie [k] v
-insertTrie k v MerklePatriciaTrie {rootNode} =
-  let newRoot = insertNode k v rootNode
-   in MerklePatriciaTrie {rootHash = updateRootHash newRoot, rootNode = newRoot}
-
--- | Generate a Merkle proof for a key
-merkleProof :: forall k v. (Eq k, Show k, Show v) => [k] -> PatriciaTrie [k] v -> Maybe (v, Digest SHA256)
-merkleProof [] Node {value} = (\v -> (v, computeNodeHash @[k] @v Node {prefix = [], value = Just v, children = []})) <$> value
-merkleProof key Node {prefix, children}
-  | prefix `isPrefixOf` key =
-      let remaining = drop (length prefix) key
-       in foldr (\child acc -> acc <|> merkleProof remaining child) Nothing children
-  | otherwise = Nothing
+    -- TODO don't like the fact, that there are 2 cases with undefined here, although logically I know, that they will not show up
+    -- compiler doesn't know that, and I'm not sure how to ensure everything is okay without some ugliness
+    -- although I can just wrap this whole acc in Maybe and deal with it. Anyway, need more thoughts
+    go :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => (Digest SHA256, [Digest SHA256]) -> Trie k a -> (Digest SHA256, [Digest SHA256])
+    go _ Empty = undefined -- Can't reach here, since 'focus' ensures that we're working with something(or it'll short cirtcuit and return Nothing)
+    go _ Leaf {} = undefined -- Can't reach here, since we start from the bottom-most branch
+    go (topHash, proof) b@Branch {} = let h = computeNodeHash b in (computeHash (topHash, h), h : proof)
diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index 69ee3b7..3c40564 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,9 +1,10 @@
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.Trie (Trie, TrieKey, lookup, insertWith, insert, empty) where
+module Data.Trie (Trie' (..), Trie, TrieKey (..), lookup, insertWith, insert, singleton, empty, module X) where
 
 import Data.Trie.Internal
+import Data.Trie.Zipper as X -- importing like this to re-export
 import Prelude hiding (lookup)
 
 -- TODO do we want to unwrap value here or return it wrapped? Maybe 2 different functions for these cases?
@@ -11,8 +12,8 @@ lookup :: (TrieKey k) => k -> Trie k a -> Maybe a
 lookup _ Empty = Nothing
 lookup k Leaf {..} = if k == key then Just value else Nothing
 lookup k Branch {..}
-  | not (matchPrefix k prefix (heightToBBit branchingBit)) = Nothing
-  | zeroBit k (heightToBBit branchingBit) = lookup k left
+  | not (matchPrefix k prefix (heightToBBit height)) = Nothing
+  | zeroBit k (heightToBBit height) = lookup k left
   | otherwise = lookup k right
 
 insertWith :: forall k a. (TrieKey k) => (a -> a -> a) -> k -> a -> Trie k a -> Trie k a
@@ -25,15 +26,18 @@ insertWith resolve k v t = go t
         then Leaf k $ resolve v value
         else join k (Leaf k v) key t
     go Branch {..} =
-      if matchPrefix k prefix (heightToBBit branchingBit)
+      if matchPrefix k prefix (heightToBBit height)
         then
-          if zeroBit k (heightToBBit branchingBit)
-            then Branch branchingBit prefix (go left) right
-            else Branch branchingBit prefix left (go right)
+          if zeroBit k (heightToBBit height)
+            then Branch height prefix (go left) right
+            else Branch height prefix left (go right)
         else join k (Leaf k v) prefix t
 
 insert :: (TrieKey k) => k -> v -> Trie k v -> Trie k v
 insert = insertWith const
 
-empty :: (TrieKey k) => Trie k v
+singleton :: (TrieKey k) => k -> a -> Trie k a
+singleton k v = insert k v empty
+
+empty :: (TrieKey k) => Trie k a
 empty = Empty
diff --git a/merkle-patricia-trie/lib/Data/Trie/Internal.hs b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
index 38bc72e..cc0ebe9 100644
--- a/merkle-patricia-trie/lib/Data/Trie/Internal.hs
+++ b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
@@ -1,26 +1,17 @@
-{-# LANGUAGE DeriveFunctor #-}
-{-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE OverloadedRecordDot #-}
-{-# LANGUAGE RecordWildCards #-}
-{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE NoImplicitPrelude #-}
 
 module Data.Trie.Internal where
 
 import Data.Bits
 import Data.Kind (Type)
-import Data.WideWord (Word256)
-import Data.Word (Word8)
 import Prelude hiding (lookup)
 
--- 0000010
-
 data Trie' k h a
   = Empty
   | Leaf {key :: k, value :: a}
   | -- | Branch node stores branching bit `bBit` and longest common prefix `pref`
-    Branch {branchingBit :: h, prefix :: k, left :: Trie' k h a, right :: Trie' k h a}
+    Branch {height :: h, prefix :: k, left :: Trie' k h a, right :: Trie' k h a}
   deriving (Show, Eq, Functor)
 
 -- | Constraint for Trie key and height. Includes default implementation for little endian Tries
diff --git a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
index ab11393..4913ad0 100644
--- a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
+++ b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
@@ -1,68 +1,111 @@
-{-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ViewPatterns #-}
 
-module Data.Trie.Zipper (TrieZipper, left, right, up, focus, zipper) where
+module Data.Trie.Zipper
+  ( TrieZipper,
+    Where (..),
+    goLeft,
+    goRight,
+    goUp,
+    focus,
+    upmostWith,
+    upmost,
+    modify,
+    zipper,
+    -- some accessors
+    getValue,
+    getKey,
+    getKeyValue,
+    getNode,
+    direction,
+  )
+where
 
-import Data.Kind (Type)
-import Data.Maybe (fromMaybe)
 import Data.Trie.Internal
 
-data Cxt (f :: Type -> Type) k h a
+data Cxt k h a
   = Top
-  | L k h (f a) (Cxt f k h a) (Trie k a)
-  | R k h (f a) (Trie k a) (Cxt f k h a)
+  | L h k (Cxt k h a) (Trie k a)
+  | R h k (Trie k a) (Cxt k h a)
   deriving (Functor)
 
-newtype TrieZipper' (f :: Type -> Type) k h a = TrieZipper {getZipper :: (Trie k a, Cxt f k h a)} deriving (Functor)
+newtype TrieZipper' k h a = TrieZipper {getZipper :: (Trie k a, Cxt k h a)} deriving (Functor)
 
-type TrieZipper (f :: Type -> Type) k a = TrieZipper' f k (TrieHeight k) a
+type TrieZipper k a = TrieZipper' k (TrieHeight k) a
 
-zipper :: (TrieKey k, Functor f) => Trie k a -> TrieZipper f k a
+zipper :: (TrieKey k) => Trie k a -> TrieZipper k a
 zipper t = TrieZipper (t, Top)
 
+data Where = LeftBranch | RightBranch deriving (Show, Eq)
+
+direction :: (TrieKey k) => TrieZipper k a -> Maybe Where
+direction (getZipper -> (_, Top)) = Nothing
+direction (getZipper -> (_, L {})) = Just LeftBranch
+direction (getZipper -> (_, R {})) = Just RightBranch
+
 -- TODO do we want to return Maybe zipper or just return the same one if the move is errorneous?
 -- Maybe 2 family of functions for this?
-left, right, up :: (TrieKey k, Applicative f) => TrieZipper f k a -> Maybe (TrieZipper f k a)
+goLeft, goRight :: (TrieKey k) => TrieZipper k a -> Maybe (TrieZipper k a)
 -- nowhere to go if we're at the leaf or in an empty tree
-left (getZipper -> (Empty, _)) = Nothing
-left (getZipper -> (Leaf {}, _)) = Nothing
--- TODO how to calculate values here? For leaf it's easy, as for branch? It must depend on the leafes nearby, I guess
-left (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (lelft, L branchingBit prefix undefined c right)
+goLeft (getZipper -> (Empty, _)) = Nothing
+goLeft (getZipper -> (Leaf {}, _)) = Nothing
+goLeft (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (left, L height prefix c right)
 -- nowhere to go if we're at the leaf or in an empty tree
-right (getZipper -> (Empty, _)) = Nothing
-right (getZipper -> (Leaf _ _, _)) = Nothing
-right (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (right, R branchingBit prefex undefined left c)
+goRight (getZipper -> (Empty, _)) = Nothing
+goRight (getZipper -> (Leaf _ _, _)) = Nothing
+goRight (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (right, R height prefix left c)
+
 -- we're already at the top
-up (getZipper -> (_, Top)) = Nothing
-up (getZipper -> (t, L branchingBit prefix v c r)) = Just $ TrieZipper (Branch branchingBit prefix t r, c)
-up (getZipper -> (t, R branchingBit prefix v l c)) = Just $ TrieZipper (Branch branchingBit prefix l t, c)
+goUp :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
+goUp z@(getZipper -> (_, Top)) = z
+goUp (getZipper -> (t, L height prefix c r)) = TrieZipper (Branch height prefix t r, c)
+goUp (getZipper -> (t, R height prefix l c)) = TrieZipper (Branch height prefix l t, c)
+
+-- | Basically a fold from the current node up to the top of the trie.
+-- | If current node is not a leaf - returns Nothing
+-- | Accepts accumulator 'b' and combine function '(b -> Trie k a -> b)', where
+-- | 'b' - starting value, 'Trie k a' - current node
+-- TODO look into how to do this through recustion schemes, since this is
+-- basically recursion scheme style recursion, we're "visiting" each node and doing something with it
+-- then accumulating the result.
+upmostWith :: (TrieKey k) => (b -> Trie k a -> b) -> b -> TrieZipper k a -> (b, TrieZipper k a)
+upmostWith _ acc z@(getZipper -> (_, Top)) = (acc, z)
+upmostWith f acc z = upmostWith f (f acc (getNode z)) $ goUp z
 
-modify :: (Applicative f) => (Trie k a -> Trie k a) -> TrieZipper f k a -> TrieZipper f k a
+upmost :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
+-- undefined should not cause any troubles due to laziness
+-- and if it does - we've made a mistake somewhere :P
+upmost = snd . upmostWith const undefined
+
+modify :: (Trie k a -> Trie k a) -> TrieZipper k a -> TrieZipper k a
 modify f (getZipper -> (t, c)) = TrieZipper (f t, c)
 
--- | Move focus to the top of the trie applying `(k -> a -> a)` along the way
--- Supplied function '(k -> a -> a)' takes key 'k' and value 'a' and returns 'a'
-unfocusWith :: (TrieKey k, Applicative f) => (k -> a -> b) -> TrieZipper f k a -> TrieZipper f k b
-unfocusWith f z@(getZipper -> (_, Top)) = fmap f z
-unfocusWith f z@(getZipper -> (Empty, _)) = fmap f z
-unfocusWith f z@(getZipper -> (Leaf key fa, c)) = unfocusWith f . fromMaybe z . up $ TrieZipper (Leaf key $ fmap (f key) fa, fmap f c)
-unfocusWith f z@(getZipper -> (Branch bBit pref l r, c)) = unfocusWith f . fromMaybe z $ up z
+getNode :: (TrieKey k) => TrieZipper k a -> Trie k a
+getNode (getZipper -> (t, _)) = t
+
+getKeyValue :: (TrieKey k) => TrieZipper k a -> Maybe (k, a)
+getKeyValue t = (,) <$> getKey t <*> getValue t
+
+getValue :: (TrieKey k) => TrieZipper k a -> Maybe a
+getValue (getZipper -> (Empty, _)) = Nothing
+getValue (getZipper -> (Branch {}, _)) = Nothing
+getValue (getZipper -> (Leaf {..}, _)) = Just value
 
--- | Move focus to the top of the trie
-unfocus :: (TrieKey k, Applicative f) => TrieZipper f k a -> TrieZipper f k a
-unfocus = upmostWith (const id)
+getKey :: (TrieKey k) => TrieZipper k a -> Maybe k
+getKey (getZipper -> (Empty, _)) = Nothing
+getKey (getZipper -> (Branch {}, _)) = Nothing
+getKey (getZipper -> (Leaf {..}, _)) = Just key
 
 -- | Create a `TrieZipper` focused on given key `k` in a trie `Trie w k a`
-focus :: (TrieKey k, Applicative f) => k -> Trie k a -> Maybe (TrieZipper f k a)
+focus :: (TrieKey k) => k -> Trie k a -> Maybe (TrieZipper k a)
 focus _ Empty = Nothing
-focus key1 t = go key1 $ TrieZipper (t, Top)
+focus key1 t = go key1 $ zipper t
   where
-    go :: (TrieKey k, Functor f) => k -> TrieZipper f k a -> Maybe (TrieZipper f k a)
+    go :: (TrieKey k) => k -> TrieZipper k a -> Maybe (TrieZipper k a)
     go _ (getZipper -> (Empty, _)) = Nothing
     go k z@(getZipper -> (Leaf {..}, _)) = if key == k then Just z else Nothing
     go k z@(getZipper -> (Branch {..}, _))
-      | not (matchPrefix k prefix (heightToBBit branchingBit)) = Nothing
-      | zeroBit k (heightToBBit branchingBit) = left z >>= go k
-      | otherwise = right z >>= go k
+      | not (matchPrefix k prefix (heightToBBit height)) = Nothing
+      | zeroBit k (heightToBBit height) = goLeft z >>= go k
+      | otherwise = goRight z >>= go k
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index 54b2e0d..a7ffc7c 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -68,8 +68,13 @@ library
     -- other-modules:
 
     -- LANGUAGE extensions used by modules in this package.
-    other-extensions:
-        NoImplicitPrelude
+    default-extensions:
+        GeneralizedNewtypeDeriving
+        OverloadedRecordDot
+        RecordWildCards
+        ScopedTypeVariables
+        DeriveFunctor
+        FlexibleContexts
 
     -- Other library packages from which modules are imported.
     build-depends:

From 393e241037443c703ad2e2a7f07e14fd329f1b4d Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Thu, 2 Jan 2025 16:54:00 +0200
Subject: [PATCH 25/31] feat: transition to Fixpoint representation; rework for
 big endian encoding[wip]

---
 .../lib/Data/Internal/RecursionSchemes.hs     |  76 ++++++++++++
 .../lib/Data/Internal/Trie.hs                 |  65 ++++++++++
 .../lib/Data/Internal/Zipper.hs               | 111 ++++++++++++++++++
 merkle-patricia-trie/lib/Data/MerkleTrie.hs   |  60 +++++-----
 merkle-patricia-trie/lib/Data/Trie.hs         |  54 +++++----
 .../lib/Data/Trie/Internal.hs                 |  61 ----------
 merkle-patricia-trie/lib/Data/Trie/Zipper.hs  | 111 ------------------
 .../merkle-patricia-trie.cabal                |  18 ++-
 merkle-patricia-trie/test/Spec.hs             |   1 +
 merkle-patricia-trie/test/TrieKeySpec.hs      |  53 +++++++++
 merkle-patricia-trie/test/TrieSpec.hs         |  61 ++++++++++
 11 files changed, 445 insertions(+), 226 deletions(-)
 create mode 100644 merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
 create mode 100644 merkle-patricia-trie/lib/Data/Internal/Trie.hs
 create mode 100644 merkle-patricia-trie/lib/Data/Internal/Zipper.hs
 delete mode 100644 merkle-patricia-trie/lib/Data/Trie/Internal.hs
 delete mode 100644 merkle-patricia-trie/lib/Data/Trie/Zipper.hs
 create mode 100644 merkle-patricia-trie/test/Spec.hs
 create mode 100644 merkle-patricia-trie/test/TrieKeySpec.hs
 create mode 100644 merkle-patricia-trie/test/TrieSpec.hs

diff --git a/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs b/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
new file mode 100644
index 0000000..35c3478
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
@@ -0,0 +1,76 @@
+module Data.Internal.RecursionSchemes where
+
+import Control.Arrow ((&&&), (<<<), (>>>))
+
+-- | Y-combinator or fixed point combinator for types
+newtype Term f = In {out :: f (Term f)}
+
+-- | Type of functions from a container `f a` to the collapsed value `a`.
+-- | Basically it's a type of functions for cata-morphisms a.k.a. folds.
+type Algebra f a = f a -> a
+
+-- | Type of functions from a value `a` to the unfolded container `f a`
+-- | Basically it's a type of functions for ana-morphisms a.k.a. unfolds.
+type CoAlgebra f a = a -> f a
+
+-- | Type of functions for folds with extra context of `Term f` we're currently working on,
+-- | and value `a` to the collapsed value `a` a.k.a. para-morphisms.
+type RAlgebra f a = f (Term f, a) -> a
+
+-- | Type of functions for unfolds that can be either terminated on `Term f` or continued on value `a`,
+-- | a.k.a. apo-morphism
+type RCoAlgebra f a = a -> f (Either (Term f) a)
+
+-- | Attributeted version of `Term f`. Used to hold the "history" of the fold that is currently going on.
+-- | Helper data type for the histo-morphism
+data Attr f a = Attr
+  { attribute :: a,
+    hole :: f (Attr f a)
+  }
+
+-- | Type of functions for folds that remembers their "history" as they go through the recursion,
+-- | a.k.a. histo-morphism
+type CVAlgebra f a = f (Attr f a) -> a
+
+-- | The dual of Attr for futu-morphisms, a.k.a. unfolds with control flow
+data CoAttr f a
+  = Automatic a
+  | Manual (f (CoAttr f a))
+
+-- | Type of functions for unfold that has control over the flow of the unfold, a.k.a. futu-morphism
+type CVCoAlgebra f a = a -> f (CoAttr f a)
+
+-- | Histomorphism
+histo :: forall f a. (Functor f) => CVAlgebra f a -> Term f -> a
+histo h = attribute . go
+  where
+    go :: Term f -> Attr f a
+    go = uncurry Attr . (h &&& id) . fmap go . out
+
+-- | Catamorphism
+cata :: (Functor f) => Algebra f a -> Term f -> a
+cata f = histo (f . fmap attribute)
+
+-- | Paramorphism
+para :: (Functor f) => RAlgebra f a -> Term f -> a
+para f = histo (f . fmap go)
+  where
+    go (Attr a h) = (In (fmap (fst . go) h), a)
+
+-- | Futumorphism
+futu :: forall f a. (Functor f) => CVCoAlgebra f a -> a -> Term f
+futu f = In . fmap go . f
+  where
+    go :: CoAttr f a -> Term f
+    go (Automatic a) = futu f a -- continue through this level
+    go (Manual g) = In (fmap go g) -- omit folding this level, delegating to the worker
+
+-- | Anamorphism
+ana :: (Functor f) => CoAlgebra f a -> a -> Term f
+ana f = futu (fmap Automatic . f)
+
+-- | Apomorphism
+apo :: (Functor f) => RCoAlgebra f a -> a -> Term f
+apo f = futu (fmap (either termToCoAttr Automatic) . f)
+  where
+    termToCoAttr = Manual . fmap termToCoAttr . out
diff --git a/merkle-patricia-trie/lib/Data/Internal/Trie.hs b/merkle-patricia-trie/lib/Data/Internal/Trie.hs
new file mode 100644
index 0000000..758eed5
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Internal/Trie.hs
@@ -0,0 +1,65 @@
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+module Data.Internal.Trie where
+
+import Data.Bits
+import Data.Internal.RecursionSchemes
+import Data.Kind (Type)
+import Prelude hiding (lookup)
+
+data TrieF k h v a
+  = Empty
+  | Leaf {key :: k, value :: v}
+  | -- | Branch node stores branching bit `bBit` and longest common prefix `pref`
+    Branch {height :: h, prefix :: k, left :: a, right :: a}
+  deriving (Show, Eq, Functor)
+
+-- | Constraint for Trie key and height. Includes default implementation for big endian Tries
+class (FiniteBits k, Bits k, Num k, Integral k, Integral (TrieHeight k), Bits (TrieHeight k), Num (TrieHeight k)) => TrieKey k where
+  -- | This will allow us to enforce the height of the trie depending on it's key
+  type TrieHeight k :: Type
+
+  -- | Discards bits not in the mask
+  mask :: k -> k -> k
+  mask k m = (k .|. (m - 1)) .&. (-m)
+
+  -- | Finds the first bit on which prefixes disagree.
+  commonBranchingBit :: k -> k -> k
+  commonBranchingBit p1 p2 = highestBit (p1 `xor` p2)
+    where
+      highestBit :: k -> k
+      highestBit x = x .&. complement (x - 1)
+
+  -- | Masks key using supplied branching bit and compares to prefix
+  matchPrefix :: k -> k -> k -> Bool
+  matchPrefix key prefix bBit = mask key bBit == prefix
+
+  -- | Converts height 'TrieHeight k' into the branching bit 'k'
+  heightToBBit :: TrieHeight k -> k
+  heightToBBit h = 2 ^ (finiteBitSize (undefined :: k) - fromIntegral h - 1)
+
+  -- | Converts branching bit 'k' into height 'TrieHeight k'
+  bBitToHeight :: k -> TrieHeight k
+  bBitToHeight = fromIntegral . (finiteBitSize (undefined :: k) - 1 -) . floor . logBase (2.0 :: Double) . fromIntegral
+
+  -- | Tests whether the desired bit is zero
+  zeroBit :: k -> k -> Bool
+  zeroBit k m = not (k `testBit` (fromIntegral (bBitToHeight m)))
+
+-- | Type alias for Trie
+type Trie k v = Term (TrieF k (TrieHeight k) v)
+
+-- | Smart constructor for Branch node
+branch :: (TrieKey k, h ~ TrieHeight k) => k -> k -> a -> a -> TrieF k h v a
+branch bBit = Branch (bBitToHeight bBit)
+
+-- | Decides how to join two tries based on prefixes(i.e. either as left or right child)
+join :: (TrieKey k) => k -> Trie k v -> k -> Trie k v -> Trie k v
+join p1 t1 p2 t2 =
+  if zeroBit p1 bBit
+    then In $ branch bBit (mask p1 bBit) t1 t2
+    else In $ branch bBit (mask p1 bBit) t2 t1
+  where
+    bBit = commonBranchingBit p1 p2
diff --git a/merkle-patricia-trie/lib/Data/Internal/Zipper.hs b/merkle-patricia-trie/lib/Data/Internal/Zipper.hs
new file mode 100644
index 0000000..8fc4a90
--- /dev/null
+++ b/merkle-patricia-trie/lib/Data/Internal/Zipper.hs
@@ -0,0 +1,111 @@
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE ViewPatterns #-}
+
+module Data.Internal.Zipper where
+
+-- ( TrieZipper,
+--   Where (..),
+--   goLeft,
+--   goRight,
+--   goUp,
+--   focus,
+--   upmostWith,
+--   upmost,
+--   modify,
+--   zipper,
+--   -- some accessors
+--   getValue,
+--   getKey,
+--   getKeyValue,
+--   getNode,
+--   direction,
+-- )
+
+-- import Data.Internal.Trie
+
+-- data Cxt k h v
+--   = Top
+--   | L h k (Cxt k h v) (Trie k v)
+--   | R h k (Trie k v) (Cxt k h v)
+--   deriving (Functor)
+
+-- newtype TrieZipper' k h a = TrieZipper {getZipper :: (Trie k a, Cxt k h a)} deriving (Functor)
+
+-- type TrieZipper k a = TrieZipper' k (TrieHeight k) a
+
+-- zipper :: (TrieKey k) => Trie k a -> TrieZipper k a
+-- zipper t = TrieZipper (t, Top)
+
+-- data Where = LeftBranch | RightBranch deriving (Show, Eq)
+
+-- direction :: (TrieKey k) => TrieZipper k a -> Maybe Where
+-- direction (getZipper -> (_, Top)) = Nothing
+-- direction (getZipper -> (_, L {})) = Just LeftBranch
+-- direction (getZipper -> (_, R {})) = Just RightBranch
+
+-- -- TODO do we want to return Maybe zipper or just return the same one if the move is errorneous?
+-- -- Maybe 2 family of functions for this?
+-- goLeft, goRight :: (TrieKey k) => TrieZipper k a -> Maybe (TrieZipper k a)
+-- -- nowhere to go if we're at the leaf or in an empty tree
+-- goLeft (getZipper -> (Empty, _)) = Nothing
+-- goLeft (getZipper -> (Leaf {}, _)) = Nothing
+-- goLeft (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (left, L height prefix c right)
+-- -- nowhere to go if we're at the leaf or in an empty tree
+-- goRight (getZipper -> (Empty, _)) = Nothing
+-- goRight (getZipper -> (Leaf _ _, _)) = Nothing
+-- goRight (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (right, R height prefix left c)
+
+-- -- we're already at the top
+-- goUp :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
+-- goUp z@(getZipper -> (_, Top)) = z
+-- goUp (getZipper -> (t, L height prefix c r)) = TrieZipper (Branch height prefix t r, c)
+-- goUp (getZipper -> (t, R height prefix l c)) = TrieZipper (Branch height prefix l t, c)
+
+-- -- | Basically a fold from the current node up to the top of the trie.
+-- -- | If current node is not a leaf - returns Nothing
+-- -- | Accepts accumulator 'b' and combine function '(b -> Trie k a -> b)', where
+-- -- | 'b' - starting value, 'Trie k a' - current node
+-- -- TODO look into how to do this through recustion schemes, since this is
+-- -- basically recursion scheme style recursion, we're "visiting" each node and doing something with it
+-- -- then accumulating the result.
+-- upmostWith :: (TrieKey k) => (b -> Trie k a -> b) -> b -> TrieZipper k a -> (b, TrieZipper k a)
+-- upmostWith _ acc z@(getZipper -> (_, Top)) = (acc, z)
+-- upmostWith f acc z = upmostWith f (f acc (getNode z)) $ goUp z
+
+-- upmost :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
+-- -- undefined should not cause any troubles due to laziness
+-- -- and if it does - we've made a mistake somewhere :P
+-- upmost = snd . upmostWith const undefined
+
+-- modify :: (Trie k a -> Trie k a) -> TrieZipper k a -> TrieZipper k a
+-- modify f (getZipper -> (t, c)) = TrieZipper (f t, c)
+
+-- getNode :: (TrieKey k) => TrieZipper k a -> Trie k a
+-- getNode (getZipper -> (t, _)) = t
+
+-- getKeyValue :: (TrieKey k) => TrieZipper k a -> Maybe (k, a)
+-- getKeyValue t = (,) <$> getKey t <*> getValue t
+
+-- getValue :: (TrieKey k) => TrieZipper k a -> Maybe a
+-- getValue (getZipper -> (Empty, _)) = Nothing
+-- getValue (getZipper -> (Branch {}, _)) = Nothing
+-- getValue (getZipper -> (Leaf {..}, _)) = Just value
+
+-- getKey :: (TrieKey k) => TrieZipper k a -> Maybe k
+-- getKey (getZipper -> (Empty, _)) = Nothing
+-- getKey (getZipper -> (Branch {}, _)) = Nothing
+-- getKey (getZipper -> (Leaf {..}, _)) = Just key
+
+-- -- | Create a `TrieZipper` focused on given key `k` in a trie `Trie w k a`
+-- focus :: (TrieKey k) => k -> Trie k a -> Maybe (TrieZipper k a)
+-- focus _ Empty = Nothing
+-- focus key1 t = go key1 $ zipper t
+--   where
+--     go :: (TrieKey k) => k -> TrieZipper k a -> Maybe (TrieZipper k a)
+--     go _ (getZipper -> (Empty, _)) = Nothing
+--     go k z@(getZipper -> (Leaf {..}, _)) = if key == k then Just z else Nothing
+--     go k z@(getZipper -> (Branch {..}, _))
+--       | not (matchPrefix k prefix (heightToBBit height)) = Nothing
+--       | zeroBit k (heightToBBit height) = goLeft z >>= go k
+--       | otherwise = goRight z >>= go k
diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index 010c762..8fc46cd 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -3,33 +3,39 @@
 
 module Data.MerkleTrie where
 
-import Crypto.Hash (Digest, SHA256, hash)
-import Data.ByteString.Char8 (pack)
-import Data.Trie
+-- import Crypto.Hash (Digest, SHA256, hash)
+-- import Data.ByteString.Char8 (pack)
+-- import Data.Trie
 
-computeHash :: (Show a) => a -> Digest SHA256
-computeHash = hash . pack . show
+-- computeHash :: (Show a) => a -> Digest SHA256
+-- computeHash = hash . pack . show
 
--- | Compute the hash of a Patricia Trie
--- | For internal use only, since it's partially matched against the trie!
--- TODO don't like the fact, that this is an incomplete pattern, although I can wrap this in Maybe
--- need to think about this
-computeNodeHash :: (Show k, Show (TrieHeight k), Show a) => Trie k a -> Digest SHA256
-computeNodeHash Leaf {..} = computeHash (key, value)
-computeNodeHash Branch {..} = computeHash (computeNodeHash left, computeNodeHash right)
+-- -- | Compute the hash of a Patricia Trie
+-- -- | For internal use only, since it's partially matched against the trie!
+-- -- TODO don't like the fact, that this is an incomplete pattern, although I can wrap this in Maybe
+-- -- need to think about this
+-- computeNodeHash :: (Show k, Show (TrieHeight k), Show a) => Trie k a -> Digest SHA256
+-- computeNodeHash Leaf {..} = computeHash (key, value)
+-- computeNodeHash Branch {..} = computeHash (computeNodeHash left, computeNodeHash right)
 
-prove :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => k -> Trie k a -> Maybe (Digest SHA256, [Digest SHA256])
-prove k t = do
-  z <- focus k t
-  startingHash <- computeHash <$> getKeyValue z -- compute hash for the leaf
-  let accum = (startingHash, []) -- initial accumulator
-      zUp = goUp z -- now we're at the bottom-most branch
-  pure . fst $ upmostWith go accum zUp
-  where
-    -- TODO don't like the fact, that there are 2 cases with undefined here, although logically I know, that they will not show up
-    -- compiler doesn't know that, and I'm not sure how to ensure everything is okay without some ugliness
-    -- although I can just wrap this whole acc in Maybe and deal with it. Anyway, need more thoughts
-    go :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => (Digest SHA256, [Digest SHA256]) -> Trie k a -> (Digest SHA256, [Digest SHA256])
-    go _ Empty = undefined -- Can't reach here, since 'focus' ensures that we're working with something(or it'll short cirtcuit and return Nothing)
-    go _ Leaf {} = undefined -- Can't reach here, since we start from the bottom-most branch
-    go (topHash, proof) b@Branch {} = let h = computeNodeHash b in (computeHash (topHash, h), h : proof)
+-- type MerkleProof = (Digest SHA256, [Digest SHA256])
+
+-- prove :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => k -> Trie k a -> Maybe MerkleProof
+-- prove k t = do
+--   z <- focus k t
+--   startingHash <- computeHash <$> getKeyValue z -- compute hash for the leaf
+--   let accum = (startingHash, []) -- initial accumulator
+--       zUp = goUp z -- now we're at the bottom-most branch
+--   pure . fst $ upmostWith go accum zUp
+--   where
+--     -- TODO I don't like the fact, that there are 2 cases with undefined here, although logically I know, that they will never show up
+--     -- compiler doesn't know that, and I'm not sure how to ensure everything is okay without some ugliness
+--     -- although I can just wrap this whole acc in Maybe and deal with it. Anyway, need more thoughts
+--     go :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => MerkleProof -> Trie k a -> MerkleProof
+--     go _ Empty = undefined -- Can't reach here, since 'focus' ensures that we're working with something(or it'll short cirtcuit and return Nothing)
+--     go _ Leaf {} = undefined -- Can't reach here, since we start from the bottom-most branch
+--     go (topHash, proof) b@Branch {} = let h = computeNodeHash b in (computeHash (topHash, h), h : proof)
+
+-- -- | TODO do this
+-- validate :: (TrieKey k) => Trie k a -> MerkleProof -> Bool
+-- validate = undefined
diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index 3c40564..25513a4 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,43 +1,47 @@
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.Trie (Trie' (..), Trie, TrieKey (..), lookup, insertWith, insert, singleton, empty, module X) where
+module Data.Trie (TrieF (..), Trie, TrieKey (..), lookup, insertWith, insert, singleton, empty, module X) where
 
-import Data.Trie.Internal
-import Data.Trie.Zipper as X -- importing like this to re-export
+-- importing like this to re-export
+
+import Data.Internal.RecursionSchemes as X
+import Data.Internal.Trie
+-- import Data.Internal.Zipper as X
 import Prelude hiding (lookup)
 
--- TODO do we want to unwrap value here or return it wrapped? Maybe 2 different functions for these cases?
-lookup :: (TrieKey k) => k -> Trie k a -> Maybe a
-lookup _ Empty = Nothing
-lookup k Leaf {..} = if k == key then Just value else Nothing
-lookup k Branch {..}
+lookup :: (TrieKey k) => k -> Trie k v -> Maybe v
+lookup _ (In Empty) = Nothing
+lookup k (In Leaf {..}) = if k == key then Just value else Nothing
+lookup k (In Branch {..})
   | not (matchPrefix k prefix (heightToBBit height)) = Nothing
-  | zeroBit k (heightToBBit height) = lookup k left
+  | k <= heightToBBit height = lookup k left
   | otherwise = lookup k right
 
-insertWith :: forall k a. (TrieKey k) => (a -> a -> a) -> k -> a -> Trie k a -> Trie k a
+-- | Inserts new value in the trie resolving possible conflicts using `(v -> v -> v)`.
+-- | Values are supplied in the following order: new value -> old value.
+insertWith :: forall k v. (TrieKey k) => (v -> v -> v) -> k -> v -> Trie k v -> Trie k v
 insertWith resolve k v t = go t
   where
-    go :: Trie k a -> Trie k a
-    go Empty = Leaf k v
-    go Leaf {..} =
+    go :: Trie k v -> Trie k v
+    go (In Empty) = In $ Leaf k v
+    go (In Leaf {..}) =
       if k == key
-        then Leaf k $ resolve v value
-        else join k (Leaf k v) key t
-    go Branch {..} =
-      if matchPrefix k prefix (heightToBBit height)
+        then In $ Leaf k $ resolve v value
+        else join k (In $ Leaf k v) key t
+    go (In Branch {..}) =
+      if matchPrefix k prefix $ heightToBBit height
         then
-          if zeroBit k (heightToBBit height)
-            then Branch height prefix (go left) right
-            else Branch height prefix left (go right)
-        else join k (Leaf k v) prefix t
+          if k <= heightToBBit height
+            then In $ Branch height prefix (go left) right
+            else In $ Branch height prefix left (go right)
+        else join k (In $ Leaf k v) prefix t
 
-insert :: (TrieKey k) => k -> v -> Trie k v -> Trie k v
+insert :: forall k v. (TrieKey k) => k -> v -> Trie k v -> Trie k v
 insert = insertWith const
 
-singleton :: (TrieKey k) => k -> a -> Trie k a
+singleton :: forall k v. (TrieKey k) => k -> v -> Trie k v
 singleton k v = insert k v empty
 
-empty :: (TrieKey k) => Trie k a
-empty = Empty
+empty :: forall k v. (TrieKey k) => Trie k v
+empty = In Empty
diff --git a/merkle-patricia-trie/lib/Data/Trie/Internal.hs b/merkle-patricia-trie/lib/Data/Trie/Internal.hs
deleted file mode 100644
index cc0ebe9..0000000
--- a/merkle-patricia-trie/lib/Data/Trie/Internal.hs
+++ /dev/null
@@ -1,61 +0,0 @@
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE NoImplicitPrelude #-}
-
-module Data.Trie.Internal where
-
-import Data.Bits
-import Data.Kind (Type)
-import Prelude hiding (lookup)
-
-data Trie' k h a
-  = Empty
-  | Leaf {key :: k, value :: a}
-  | -- | Branch node stores branching bit `bBit` and longest common prefix `pref`
-    Branch {height :: h, prefix :: k, left :: Trie' k h a, right :: Trie' k h a}
-  deriving (Show, Eq, Functor)
-
--- | Constraint for Trie key and height. Includes default implementation for little endian Tries
--- TODO replace default with big endian Trie
-class (Bits k, Num k, Integral k, Integral (TrieHeight k)) => TrieKey k where
-  -- | This will allow us to enforce the height of the trie depending on it's key
-  type TrieHeight k :: Type
-
-  -- | Discards bits not in the mask
-  mask :: k -> k -> k
-  mask k m = k .&. (m - 1)
-
-  -- | Finds the first bit on which prefixes disagree.
-  commonBranchingBit :: k -> k -> k
-  commonBranchingBit p1 p2 = lowestBit (p1 `xor` p2)
-    where
-      lowestBit :: k -> k
-      lowestBit x = x .&. complement x
-
-  -- | Masks key using supplied branching bit and compares to prefix
-  matchPrefix :: k -> k -> k -> Bool
-  matchPrefix key prefix m = mask key m == prefix
-
-  -- | Converts height 'TrieHeight k' into the branching bit 'k'
-  heightToBBit :: TrieHeight k -> k
-  heightToBBit h = 2 ^ h
-
-  -- | Converts branching bit 'k' into height 'TrieHeight k'
-  bBitToHeight :: k -> TrieHeight k
-  bBitToHeight = floor . logBase 2.0 . fromIntegral
-
-  -- | Tests whether the desired bit is zero
-  zeroBit :: k -> k -> Bool
-  zeroBit k m = (k .&. m) == 0
-
-type Trie k a = Trie' k (TrieHeight k) a
-
-branch :: (TrieKey k) => k -> k -> Trie k a -> Trie k a -> Trie k a
-branch bBit = Branch (bBitToHeight bBit)
-
-join :: (TrieKey k) => k -> Trie k a -> k -> Trie k a -> Trie k a
-join p1 t1 p2 t2 =
-  if zeroBit p1 bBit
-    then branch bBit (mask p1 bBit) t1 t2
-    else branch bBit (mask p1 bBit) t2 t1
-  where
-    bBit = commonBranchingBit p1 p2
diff --git a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs b/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
deleted file mode 100644
index 4913ad0..0000000
--- a/merkle-patricia-trie/lib/Data/Trie/Zipper.hs
+++ /dev/null
@@ -1,111 +0,0 @@
-{-# LANGUAGE KindSignatures #-}
-{-# LANGUAGE RecordWildCards #-}
-{-# LANGUAGE ViewPatterns #-}
-
-module Data.Trie.Zipper
-  ( TrieZipper,
-    Where (..),
-    goLeft,
-    goRight,
-    goUp,
-    focus,
-    upmostWith,
-    upmost,
-    modify,
-    zipper,
-    -- some accessors
-    getValue,
-    getKey,
-    getKeyValue,
-    getNode,
-    direction,
-  )
-where
-
-import Data.Trie.Internal
-
-data Cxt k h a
-  = Top
-  | L h k (Cxt k h a) (Trie k a)
-  | R h k (Trie k a) (Cxt k h a)
-  deriving (Functor)
-
-newtype TrieZipper' k h a = TrieZipper {getZipper :: (Trie k a, Cxt k h a)} deriving (Functor)
-
-type TrieZipper k a = TrieZipper' k (TrieHeight k) a
-
-zipper :: (TrieKey k) => Trie k a -> TrieZipper k a
-zipper t = TrieZipper (t, Top)
-
-data Where = LeftBranch | RightBranch deriving (Show, Eq)
-
-direction :: (TrieKey k) => TrieZipper k a -> Maybe Where
-direction (getZipper -> (_, Top)) = Nothing
-direction (getZipper -> (_, L {})) = Just LeftBranch
-direction (getZipper -> (_, R {})) = Just RightBranch
-
--- TODO do we want to return Maybe zipper or just return the same one if the move is errorneous?
--- Maybe 2 family of functions for this?
-goLeft, goRight :: (TrieKey k) => TrieZipper k a -> Maybe (TrieZipper k a)
--- nowhere to go if we're at the leaf or in an empty tree
-goLeft (getZipper -> (Empty, _)) = Nothing
-goLeft (getZipper -> (Leaf {}, _)) = Nothing
-goLeft (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (left, L height prefix c right)
--- nowhere to go if we're at the leaf or in an empty tree
-goRight (getZipper -> (Empty, _)) = Nothing
-goRight (getZipper -> (Leaf _ _, _)) = Nothing
-goRight (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (right, R height prefix left c)
-
--- we're already at the top
-goUp :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
-goUp z@(getZipper -> (_, Top)) = z
-goUp (getZipper -> (t, L height prefix c r)) = TrieZipper (Branch height prefix t r, c)
-goUp (getZipper -> (t, R height prefix l c)) = TrieZipper (Branch height prefix l t, c)
-
--- | Basically a fold from the current node up to the top of the trie.
--- | If current node is not a leaf - returns Nothing
--- | Accepts accumulator 'b' and combine function '(b -> Trie k a -> b)', where
--- | 'b' - starting value, 'Trie k a' - current node
--- TODO look into how to do this through recustion schemes, since this is
--- basically recursion scheme style recursion, we're "visiting" each node and doing something with it
--- then accumulating the result.
-upmostWith :: (TrieKey k) => (b -> Trie k a -> b) -> b -> TrieZipper k a -> (b, TrieZipper k a)
-upmostWith _ acc z@(getZipper -> (_, Top)) = (acc, z)
-upmostWith f acc z = upmostWith f (f acc (getNode z)) $ goUp z
-
-upmost :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
--- undefined should not cause any troubles due to laziness
--- and if it does - we've made a mistake somewhere :P
-upmost = snd . upmostWith const undefined
-
-modify :: (Trie k a -> Trie k a) -> TrieZipper k a -> TrieZipper k a
-modify f (getZipper -> (t, c)) = TrieZipper (f t, c)
-
-getNode :: (TrieKey k) => TrieZipper k a -> Trie k a
-getNode (getZipper -> (t, _)) = t
-
-getKeyValue :: (TrieKey k) => TrieZipper k a -> Maybe (k, a)
-getKeyValue t = (,) <$> getKey t <*> getValue t
-
-getValue :: (TrieKey k) => TrieZipper k a -> Maybe a
-getValue (getZipper -> (Empty, _)) = Nothing
-getValue (getZipper -> (Branch {}, _)) = Nothing
-getValue (getZipper -> (Leaf {..}, _)) = Just value
-
-getKey :: (TrieKey k) => TrieZipper k a -> Maybe k
-getKey (getZipper -> (Empty, _)) = Nothing
-getKey (getZipper -> (Branch {}, _)) = Nothing
-getKey (getZipper -> (Leaf {..}, _)) = Just key
-
--- | Create a `TrieZipper` focused on given key `k` in a trie `Trie w k a`
-focus :: (TrieKey k) => k -> Trie k a -> Maybe (TrieZipper k a)
-focus _ Empty = Nothing
-focus key1 t = go key1 $ zipper t
-  where
-    go :: (TrieKey k) => k -> TrieZipper k a -> Maybe (TrieZipper k a)
-    go _ (getZipper -> (Empty, _)) = Nothing
-    go k z@(getZipper -> (Leaf {..}, _)) = if key == k then Just z else Nothing
-    go k z@(getZipper -> (Branch {..}, _))
-      | not (matchPrefix k prefix (heightToBBit height)) = Nothing
-      | zeroBit k (heightToBBit height) = goLeft z >>= go k
-      | otherwise = goRight z >>= go k
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index a7ffc7c..2fa77ef 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -60,8 +60,10 @@ library
     exposed-modules:
         Data.MerkleTrie
         Data.Trie
-        Data.Trie.Internal
-        Data.Trie.Zipper
+        -- Data.Internal
+        Data.Internal.Trie
+        Data.Internal.Zipper
+        Data.Internal.RecursionSchemes
         Data.Utils
 
     -- Modules included in this library but not exported.
@@ -90,3 +92,15 @@ library
 
     -- Base language which the package is written in.
     default-language: Haskell2010
+
+test-suite test
+    type: exitcode-stdio-1.0
+    hs-source-dirs: test
+    main-is: Spec.hs
+    build-depends:
+        base,
+        merkle-patricia-trie,
+        hspec,
+        wide-word,
+        QuickCheck
+    default-language: Haskell2010
diff --git a/merkle-patricia-trie/test/Spec.hs b/merkle-patricia-trie/test/Spec.hs
new file mode 100644
index 0000000..a824f8c
--- /dev/null
+++ b/merkle-patricia-trie/test/Spec.hs
@@ -0,0 +1 @@
+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}
diff --git a/merkle-patricia-trie/test/TrieKeySpec.hs b/merkle-patricia-trie/test/TrieKeySpec.hs
new file mode 100644
index 0000000..dd90ffe
--- /dev/null
+++ b/merkle-patricia-trie/test/TrieKeySpec.hs
@@ -0,0 +1,53 @@
+{-# LANGUAGE BinaryLiterals #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module TrieKeySpec (spec) where
+
+import Data.Bits
+import Data.Internal.Trie
+import Data.WideWord (Word256)
+import Data.Word (Word8)
+import Test.Hspec
+import Test.QuickCheck hiding ((.&.))
+
+instance TrieKey Word256 where
+  type TrieHeight Word256 = Word8
+
+instance Arbitrary Word256 where
+  arbitrary = arbitraryBoundedEnum
+
+spec :: Spec
+spec = do
+  describe "TrieKey typeclass" $ do
+    it "mask should discard bits not in the mask" $ do
+      mask @Word256 0b1101 0b0 `shouldBe` 0b0 -- zero mask should return zero
+      mask @Word256 0b1101 0b0100 `shouldBe` 0b1100
+      mask @Word256 0b1101 0b0010 `shouldBe` 0b1100
+      mask @Word256 0b1101 0b0001 `shouldBe` 0b1101
+
+    it "commonBranchingBit should find the first bit on which prefixes disagree" $ do
+      commonBranchingBit @Word256 0b1100 0b1000 `shouldBe` 0b0100
+      commonBranchingBit @Word256 0b1100 0b1110 `shouldBe` 0b0010
+
+    it "matchPrefix should mask the key using supplied branching bit and compare to prefix" $ do
+      matchPrefix @Word256 0b1101 0b1100 0b0100 `shouldBe` True
+      matchPrefix @Word256 0b1101 0b1000 0b0100 `shouldBe` False
+
+    it "heightToBBit should convert height Word256 to the branching bit" $ do
+      heightToBBit @Word256 0 `shouldBe` setBit 0 255
+      heightToBBit @Word256 1 `shouldBe` setBit 0 254
+
+    it "bBitToHeight should convert branching bit Word256 to height" $ do
+      bBitToHeight @Word256 (setBit 0 255) `shouldBe` 0
+      bBitToHeight @Word256 (setBit 0 254) `shouldBe` 1
+
+    it "zeroBit should test whether the desired bit is zero" $ do
+      zeroBit @Word256 0b1101 (heightToBBit @Word256 1) `shouldBe` True
+      zeroBit @Word256 0b1101 (heightToBBit @Word256 2) `shouldBe` False
+
+    it "QuickCheck property: heightToBBit and bBitToHeight should be inverses" $
+      property $
+        \(h :: Word8) ->
+          bBitToHeight @Word256 (heightToBBit @Word256 h) == h
diff --git a/merkle-patricia-trie/test/TrieSpec.hs b/merkle-patricia-trie/test/TrieSpec.hs
new file mode 100644
index 0000000..d7517c7
--- /dev/null
+++ b/merkle-patricia-trie/test/TrieSpec.hs
@@ -0,0 +1,61 @@
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module TrieSpec (spec) where
+
+import qualified Data.Trie as Trie
+import Data.WideWord (Word256)
+import Data.Word (Word, Word8)
+import Test.Hspec
+import Test.QuickCheck
+
+instance Trie.TrieKey Word256 where
+  type TrieHeight Word256 = Word8
+
+spec :: Spec
+spec = do
+  describe "Data.Trie" $ do
+    it "lookup should return Nothing for an empty trie" $ do
+      (Trie.lookup @Word256 @_ 1 Trie.empty :: Maybe Word256) `shouldBe` Nothing
+
+    it "lookup should return the value for a key in a singleton trie" $ do
+      let t = Trie.singleton @Word256 @_ 1 "value"
+      Trie.lookup 1 t `shouldBe` Just "value"
+
+    it "lookup should return Nothing for a key not in the trie" $ do
+      let t = Trie.singleton @Word256 @_ 1 "value"
+      Trie.lookup 2 t `shouldBe` Nothing
+
+    it "insert should add a key-value pair to an empty trie" $ do
+      let t = Trie.insert @Word256 @_ 1 "value" Trie.empty
+      Trie.lookup 1 t `shouldBe` Just "value"
+
+    it "insertWith should resolve conflicts using the provided function" $ do
+      let t = Trie.insertWith @Word256 @_ (++) 1 "value1" $ Trie.insert 1 "value2" Trie.empty
+      Trie.lookup 1 t `shouldBe` Just "value1value2"
+
+    it "insert should overwrite the value for an existing key" $ do
+      let t = Trie.insert @Word256 @_ 1 "newValue" $ Trie.insert 1 "oldValue" Trie.empty
+      Trie.lookup 1 t `shouldBe` Just "newValue"
+
+    it "singleton should create a trie with one key-value pair" $ do
+      let t = Trie.singleton @Word256 @_ 1 "value"
+      Trie.lookup 1 t `shouldBe` Just "value"
+
+    it "empty should create an empty trie" $ do
+      (Trie.lookup @Word256 @_ 1 Trie.empty :: Maybe Word256) `shouldBe` Nothing
+
+    it "insert and lookup should work for multiple keys" $ do
+      let t = Trie.insert @Word256 @_ 2 "value2" $ Trie.insert 1 "value1" Trie.empty
+      Trie.lookup 1 t `shouldBe` Just "value1"
+      Trie.lookup 2 t `shouldBe` Just "value2"
+
+    it "insertWith should handle multiple keys correctly" $ do
+      let t = Trie.insertWith @Word256 @_ (++) 2 "value2" $ Trie.insertWith (++) 1 "value1" Trie.empty
+      Trie.lookup 1 t `shouldBe` Just "value1"
+      Trie.lookup 2 t `shouldBe` Just "value2"
+
+    it "insertWith should handle conflicts correctly" $ do
+      let t = Trie.insertWith @Word256 @_ (++) 1 "value1" $ Trie.insertWith (++) 1 "value2" Trie.empty
+      Trie.lookup 1 t `shouldBe` Just "value1value2"

From aa5049f08239e85bc12cdc70b628902fd1360efe Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Thu, 2 Jan 2025 21:14:52 +0200
Subject: [PATCH 26/31] feat: fix proof and validate[wip]

---
 .../lib/Data/Internal/RecursionSchemes.hs     |  2 +-
 .../lib/Data/Internal/Trie.hs                 |  5 +-
 merkle-patricia-trie/lib/Data/MerkleTrie.hs   | 75 +++++++++++--------
 merkle-patricia-trie/lib/Data/Trie.hs         |  2 +-
 4 files changed, 50 insertions(+), 34 deletions(-)

diff --git a/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs b/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
index 35c3478..5f7eee2 100644
--- a/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
+++ b/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
@@ -1,6 +1,6 @@
 module Data.Internal.RecursionSchemes where
 
-import Control.Arrow ((&&&), (<<<), (>>>))
+import Control.Arrow ((&&&))
 
 -- | Y-combinator or fixed point combinator for types
 newtype Term f = In {out :: f (Term f)}
diff --git a/merkle-patricia-trie/lib/Data/Internal/Trie.hs b/merkle-patricia-trie/lib/Data/Internal/Trie.hs
index 758eed5..29b9aaa 100644
--- a/merkle-patricia-trie/lib/Data/Internal/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Internal/Trie.hs
@@ -48,8 +48,11 @@ class (FiniteBits k, Bits k, Num k, Integral k, Integral (TrieHeight k), Bits (T
   zeroBit :: k -> k -> Bool
   zeroBit k m = not (k `testBit` (fromIntegral (bBitToHeight m)))
 
+-- | Type alias for TrieF
+type TrieF' k v = TrieF k (TrieHeight k) v
+
 -- | Type alias for Trie
-type Trie k v = Term (TrieF k (TrieHeight k) v)
+type Trie k v = Term (TrieF' k v)
 
 -- | Smart constructor for Branch node
 branch :: (TrieKey k, h ~ TrieHeight k) => k -> k -> a -> a -> TrieF k h v a
diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index 8fc46cd..d5d8399 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -3,39 +3,52 @@
 
 module Data.MerkleTrie where
 
--- import Crypto.Hash (Digest, SHA256, hash)
--- import Data.ByteString.Char8 (pack)
--- import Data.Trie
+import Crypto.Hash (Digest, SHA256, hash)
+import Data.ByteString.Char8 (pack)
+import Data.Trie
 
--- computeHash :: (Show a) => a -> Digest SHA256
--- computeHash = hash . pack . show
+computeHash :: (Show a) => a -> Digest SHA256
+computeHash = hash . pack . show
 
--- -- | Compute the hash of a Patricia Trie
--- -- | For internal use only, since it's partially matched against the trie!
--- -- TODO don't like the fact, that this is an incomplete pattern, although I can wrap this in Maybe
--- -- need to think about this
--- computeNodeHash :: (Show k, Show (TrieHeight k), Show a) => Trie k a -> Digest SHA256
--- computeNodeHash Leaf {..} = computeHash (key, value)
--- computeNodeHash Branch {..} = computeHash (computeNodeHash left, computeNodeHash right)
+-- | Compute the hash of a Patricia Trie
+-- | For internal use only, since it's partially matched against the trie!
+-- TODO don't like the fact, that this is an incomplete pattern, although I can wrap this in Maybe
+-- need to think about this
+-- | Compute the hash of a Patricia Trie node
+computeNodeHash :: (Show k, Show (TrieHeight k), Show v) => TrieF k (TrieHeight k) v (Digest SHA256) -> Digest SHA256
+computeNodeHash Empty = computeHash "Empty"
+computeNodeHash Leaf {..} = computeHash (key, value)
+computeNodeHash Branch {..} = computeHash (height, prefix, left, right)
 
--- type MerkleProof = (Digest SHA256, [Digest SHA256])
+type MerkleProof = (Digest SHA256, [Digest SHA256])
 
--- prove :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => k -> Trie k a -> Maybe MerkleProof
--- prove k t = do
---   z <- focus k t
---   startingHash <- computeHash <$> getKeyValue z -- compute hash for the leaf
---   let accum = (startingHash, []) -- initial accumulator
---       zUp = goUp z -- now we're at the bottom-most branch
---   pure . fst $ upmostWith go accum zUp
---   where
---     -- TODO I don't like the fact, that there are 2 cases with undefined here, although logically I know, that they will never show up
---     -- compiler doesn't know that, and I'm not sure how to ensure everything is okay without some ugliness
---     -- although I can just wrap this whole acc in Maybe and deal with it. Anyway, need more thoughts
---     go :: (Show a, Show k, Show (TrieHeight k), TrieKey k) => MerkleProof -> Trie k a -> MerkleProof
---     go _ Empty = undefined -- Can't reach here, since 'focus' ensures that we're working with something(or it'll short cirtcuit and return Nothing)
---     go _ Leaf {} = undefined -- Can't reach here, since we start from the bottom-most branch
---     go (topHash, proof) b@Branch {} = let h = computeNodeHash b in (computeHash (topHash, h), h : proof)
+-- | Generate a Merkle proof for a given key in the trie
+proof :: forall k v. (Show k, Show (TrieHeight k), Show v, TrieKey k) => k -> Trie k v -> Maybe MerkleProof
+proof k t = cata go t
+  where
+    go :: Algebra (TrieF' k v) (Maybe MerkleProof)
+    go Empty = Nothing
+    -- for every leaf we just compute it's hash
+    go Leaf {..} = Just (computeHash (key, value), [])
+    -- for every branch we look at the prefix
+    go Branch {..}
+      -- if the key doesn't match the prefix, we calculate the hash and store it
+      | not (matchPrefix k prefix (heightToBBit height)) = case (left, right) of
+        (Just (hl, _), Just (hr, _)) -> Just (computeHash (height, prefix, hl, hr), [])
+        (_, _) -> Nothing
+      -- if the key matches the prefix and we're coming from the left branch(zeroBit check),
+      -- then we take the hash from there and append the hash of the current node to the path
+      | zeroBit k (heightToBBit height) = case (left, right) of
+        (Just (hl, p), Just (hr, _)) -> Just (hl, (computeHash (height, prefix, hl, hr)) : p)
+        _ -> Nothing
+      -- the same here, just for the right branch
+      | otherwise = case (left, right) of
+        (Just (hl, _), Just (hr, p)) -> Just (hr, (computeHash (height, prefix, hl, hr)) : p)
+        _ -> Nothing
 
--- -- | TODO do this
--- validate :: (TrieKey k) => Trie k a -> MerkleProof -> Bool
--- validate = undefined
+-- | Validate a Merkle proof against the root hash of the trie
+validate :: Digest SHA256 -> MerkleProof -> Bool
+validate rootHash (leafHash, pf) = foldl go leafHash pf == rootHash
+  where
+    go :: Digest SHA256 -> Digest SHA256 -> Digest SHA256
+    go acc h = computeHash (acc, h)
diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index 25513a4..5e4704b 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,7 +1,7 @@
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.Trie (TrieF (..), Trie, TrieKey (..), lookup, insertWith, insert, singleton, empty, module X) where
+module Data.Trie (TrieF (..), TrieF', Trie, TrieKey (..), lookup, insertWith, insert, singleton, empty, module X) where
 
 -- importing like this to re-export
 

From 39ee0fa6106cc3f8a1be4fa18c022857a1241d2d Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Thu, 2 Jan 2025 21:19:03 +0200
Subject: [PATCH 27/31] docs: some todos for Fare

---
 merkle-patricia-trie/lib/Data/MerkleTrie.hs | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index d5d8399..4a03fc1 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -23,6 +23,11 @@ computeNodeHash Branch {..} = computeHash (height, prefix, left, right)
 type MerkleProof = (Digest SHA256, [Digest SHA256])
 
 -- | Generate a Merkle proof for a given key in the trie
+-- TODO: check the implementation. Right now it's going through the whole trie and picks the hashes of the nodes
+-- that match the key. In the branch, when the branch is one of the branches that lead to the root, it computes
+-- the hash using the hash of the leaf and the trie on the other side. Is this right?
+-- If not, might need to use paramorphism here, as it allows to see the nodes we're operating on alongside the accumulator.
+-- That will allow us to properly calculate the hash of the trie from where we came to the current branch.
 proof :: forall k v. (Show k, Show (TrieHeight k), Show v, TrieKey k) => k -> Trie k v -> Maybe MerkleProof
 proof k t = cata go t
   where

From 6704a08dac2e55e7e290445e1f492fb29647abc4 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 6 Jan 2025 12:08:48 +0200
Subject: [PATCH 28/31] feat: merkle proof and validation[wip]

---
 .../lib/Data/Internal/RecursionSchemes.hs     |  10 +-
 .../lib/Data/Internal/Trie.hs                 |  16 +--
 .../lib/Data/Internal/Zipper.hs               | 111 ----------------
 merkle-patricia-trie/lib/Data/MerkleTrie.hs   | 125 +++++++++++-------
 merkle-patricia-trie/lib/Data/Trie.hs         |  41 +++++-
 .../merkle-patricia-trie.cabal                |   1 -
 merkle-patricia-trie/test/TrieKeySpec.hs      |  26 ++--
 7 files changed, 145 insertions(+), 185 deletions(-)
 delete mode 100644 merkle-patricia-trie/lib/Data/Internal/Zipper.hs

diff --git a/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs b/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
index 5f7eee2..e2a0c58 100644
--- a/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
+++ b/merkle-patricia-trie/lib/Data/Internal/RecursionSchemes.hs
@@ -51,10 +51,18 @@ histo h = attribute . go
 cata :: (Functor f) => Algebra f a -> Term f -> a
 cata f = histo (f . fmap attribute)
 
+-- | Helper function to transform 'Term f -> f a -> a' into 'RAlgebra f a'
+transform :: forall f a. (Functor f) => (Term f -> f a -> a) -> RAlgebra f a
+transform h fta = h term fa
+  where
+    term = In $ fmap fst fta
+    fa = fmap snd fta
+
 -- | Paramorphism
-para :: (Functor f) => RAlgebra f a -> Term f -> a
+para :: forall f a. (Functor f) => RAlgebra f a -> Term f -> a
 para f = histo (f . fmap go)
   where
+    go :: Attr f a -> (Term f, a)
     go (Attr a h) = (In (fmap (fst . go) h), a)
 
 -- | Futumorphism
diff --git a/merkle-patricia-trie/lib/Data/Internal/Trie.hs b/merkle-patricia-trie/lib/Data/Internal/Trie.hs
index 29b9aaa..2a4af81 100644
--- a/merkle-patricia-trie/lib/Data/Internal/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Internal/Trie.hs
@@ -16,21 +16,21 @@ data TrieF k h v a
     Branch {height :: h, prefix :: k, left :: a, right :: a}
   deriving (Show, Eq, Functor)
 
--- | Constraint for Trie key and height. Includes default implementation for big endian Tries
+-- | Constraint for Trie key and height. Provides default methods for calculating neccessary info.
 class (FiniteBits k, Bits k, Num k, Integral k, Integral (TrieHeight k), Bits (TrieHeight k), Num (TrieHeight k)) => TrieKey k where
   -- | This will allow us to enforce the height of the trie depending on it's key
   type TrieHeight k :: Type
 
   -- | Discards bits not in the mask
   mask :: k -> k -> k
-  mask k m = (k .|. (m - 1)) .&. (-m)
+  mask k m = k .&. (m - 1)
 
   -- | Finds the first bit on which prefixes disagree.
   commonBranchingBit :: k -> k -> k
-  commonBranchingBit p1 p2 = highestBit (p1 `xor` p2)
+  commonBranchingBit p1 p2 = lowesBit (p1 `xor` p2)
     where
-      highestBit :: k -> k
-      highestBit x = x .&. complement (x - 1)
+      lowesBit :: k -> k
+      lowesBit x = x .&. (-x)
 
   -- | Masks key using supplied branching bit and compares to prefix
   matchPrefix :: k -> k -> k -> Bool
@@ -38,15 +38,15 @@ class (FiniteBits k, Bits k, Num k, Integral k, Integral (TrieHeight k), Bits (T
 
   -- | Converts height 'TrieHeight k' into the branching bit 'k'
   heightToBBit :: TrieHeight k -> k
-  heightToBBit h = 2 ^ (finiteBitSize (undefined :: k) - fromIntegral h - 1)
+  heightToBBit = (2 ^)
 
   -- | Converts branching bit 'k' into height 'TrieHeight k'
   bBitToHeight :: k -> TrieHeight k
-  bBitToHeight = fromIntegral . (finiteBitSize (undefined :: k) - 1 -) . floor . logBase (2.0 :: Double) . fromIntegral
+  bBitToHeight = floor . logBase (2.0 :: Double) . fromIntegral
 
   -- | Tests whether the desired bit is zero
   zeroBit :: k -> k -> Bool
-  zeroBit k m = not (k `testBit` (fromIntegral (bBitToHeight m)))
+  zeroBit k m = not (k `testBit` fromIntegral (bBitToHeight m))
 
 -- | Type alias for TrieF
 type TrieF' k v = TrieF k (TrieHeight k) v
diff --git a/merkle-patricia-trie/lib/Data/Internal/Zipper.hs b/merkle-patricia-trie/lib/Data/Internal/Zipper.hs
deleted file mode 100644
index 8fc4a90..0000000
--- a/merkle-patricia-trie/lib/Data/Internal/Zipper.hs
+++ /dev/null
@@ -1,111 +0,0 @@
-{-# LANGUAGE KindSignatures #-}
-{-# LANGUAGE RecordWildCards #-}
-{-# LANGUAGE ViewPatterns #-}
-
-module Data.Internal.Zipper where
-
--- ( TrieZipper,
---   Where (..),
---   goLeft,
---   goRight,
---   goUp,
---   focus,
---   upmostWith,
---   upmost,
---   modify,
---   zipper,
---   -- some accessors
---   getValue,
---   getKey,
---   getKeyValue,
---   getNode,
---   direction,
--- )
-
--- import Data.Internal.Trie
-
--- data Cxt k h v
---   = Top
---   | L h k (Cxt k h v) (Trie k v)
---   | R h k (Trie k v) (Cxt k h v)
---   deriving (Functor)
-
--- newtype TrieZipper' k h a = TrieZipper {getZipper :: (Trie k a, Cxt k h a)} deriving (Functor)
-
--- type TrieZipper k a = TrieZipper' k (TrieHeight k) a
-
--- zipper :: (TrieKey k) => Trie k a -> TrieZipper k a
--- zipper t = TrieZipper (t, Top)
-
--- data Where = LeftBranch | RightBranch deriving (Show, Eq)
-
--- direction :: (TrieKey k) => TrieZipper k a -> Maybe Where
--- direction (getZipper -> (_, Top)) = Nothing
--- direction (getZipper -> (_, L {})) = Just LeftBranch
--- direction (getZipper -> (_, R {})) = Just RightBranch
-
--- -- TODO do we want to return Maybe zipper or just return the same one if the move is errorneous?
--- -- Maybe 2 family of functions for this?
--- goLeft, goRight :: (TrieKey k) => TrieZipper k a -> Maybe (TrieZipper k a)
--- -- nowhere to go if we're at the leaf or in an empty tree
--- goLeft (getZipper -> (Empty, _)) = Nothing
--- goLeft (getZipper -> (Leaf {}, _)) = Nothing
--- goLeft (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (left, L height prefix c right)
--- -- nowhere to go if we're at the leaf or in an empty tree
--- goRight (getZipper -> (Empty, _)) = Nothing
--- goRight (getZipper -> (Leaf _ _, _)) = Nothing
--- goRight (getZipper -> (Branch {..}, c)) = Just $ TrieZipper (right, R height prefix left c)
-
--- -- we're already at the top
--- goUp :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
--- goUp z@(getZipper -> (_, Top)) = z
--- goUp (getZipper -> (t, L height prefix c r)) = TrieZipper (Branch height prefix t r, c)
--- goUp (getZipper -> (t, R height prefix l c)) = TrieZipper (Branch height prefix l t, c)
-
--- -- | Basically a fold from the current node up to the top of the trie.
--- -- | If current node is not a leaf - returns Nothing
--- -- | Accepts accumulator 'b' and combine function '(b -> Trie k a -> b)', where
--- -- | 'b' - starting value, 'Trie k a' - current node
--- -- TODO look into how to do this through recustion schemes, since this is
--- -- basically recursion scheme style recursion, we're "visiting" each node and doing something with it
--- -- then accumulating the result.
--- upmostWith :: (TrieKey k) => (b -> Trie k a -> b) -> b -> TrieZipper k a -> (b, TrieZipper k a)
--- upmostWith _ acc z@(getZipper -> (_, Top)) = (acc, z)
--- upmostWith f acc z = upmostWith f (f acc (getNode z)) $ goUp z
-
--- upmost :: (TrieKey k) => TrieZipper k a -> TrieZipper k a
--- -- undefined should not cause any troubles due to laziness
--- -- and if it does - we've made a mistake somewhere :P
--- upmost = snd . upmostWith const undefined
-
--- modify :: (Trie k a -> Trie k a) -> TrieZipper k a -> TrieZipper k a
--- modify f (getZipper -> (t, c)) = TrieZipper (f t, c)
-
--- getNode :: (TrieKey k) => TrieZipper k a -> Trie k a
--- getNode (getZipper -> (t, _)) = t
-
--- getKeyValue :: (TrieKey k) => TrieZipper k a -> Maybe (k, a)
--- getKeyValue t = (,) <$> getKey t <*> getValue t
-
--- getValue :: (TrieKey k) => TrieZipper k a -> Maybe a
--- getValue (getZipper -> (Empty, _)) = Nothing
--- getValue (getZipper -> (Branch {}, _)) = Nothing
--- getValue (getZipper -> (Leaf {..}, _)) = Just value
-
--- getKey :: (TrieKey k) => TrieZipper k a -> Maybe k
--- getKey (getZipper -> (Empty, _)) = Nothing
--- getKey (getZipper -> (Branch {}, _)) = Nothing
--- getKey (getZipper -> (Leaf {..}, _)) = Just key
-
--- -- | Create a `TrieZipper` focused on given key `k` in a trie `Trie w k a`
--- focus :: (TrieKey k) => k -> Trie k a -> Maybe (TrieZipper k a)
--- focus _ Empty = Nothing
--- focus key1 t = go key1 $ zipper t
---   where
---     go :: (TrieKey k) => k -> TrieZipper k a -> Maybe (TrieZipper k a)
---     go _ (getZipper -> (Empty, _)) = Nothing
---     go k z@(getZipper -> (Leaf {..}, _)) = if key == k then Just z else Nothing
---     go k z@(getZipper -> (Branch {..}, _))
---       | not (matchPrefix k prefix (heightToBBit height)) = Nothing
---       | zeroBit k (heightToBBit height) = goLeft z >>= go k
---       | otherwise = goRight z >>= go k
diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index 4a03fc1..cac8c13 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -1,59 +1,94 @@
+{-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.MerkleTrie where
+module Data.MerkleTrie
+  ( proof,
+    validate,
+    computeRootHash,
+    merkelize,
+    MerkleProof (..),
+    MerkleTrie (..),
+  )
+where
 
-import Crypto.Hash (Digest, SHA256, hash)
+import Crypto.Hash (Blake2b_256, Digest, hash)
 import Data.ByteString.Char8 (pack)
 import Data.Trie
 
-computeHash :: (Show a) => a -> Digest SHA256
-computeHash = hash . pack . show
+data MerkleTrie k v = MerkleTrie
+  { rootHash :: Digest Blake2b_256,
+    trie :: Trie k v
+  }
+
+computeRootHash :: forall k v. (Show k, Show v, Show (TrieHeight k)) => Trie k v -> Maybe (Digest Blake2b_256)
+computeRootHash = cata go
+  where
+    go :: Algebra (TrieF' k v) (Maybe (Digest Blake2b_256))
+    go Empty = Nothing
+    go Leaf {..} = Just $ computeHash (key, value)
+    go Branch {..} = Just $ computeHash (height, prefix, left, right)
 
--- | Compute the hash of a Patricia Trie
--- | For internal use only, since it's partially matched against the trie!
--- TODO don't like the fact, that this is an incomplete pattern, although I can wrap this in Maybe
--- need to think about this
--- | Compute the hash of a Patricia Trie node
-computeNodeHash :: (Show k, Show (TrieHeight k), Show v) => TrieF k (TrieHeight k) v (Digest SHA256) -> Digest SHA256
-computeNodeHash Empty = computeHash "Empty"
-computeNodeHash Leaf {..} = computeHash (key, value)
-computeNodeHash Branch {..} = computeHash (height, prefix, left, right)
+merkelize :: (Show k, Show v, Show (TrieHeight k)) => Trie k v -> Maybe (MerkleTrie k v)
+merkelize trie = let rootHash = computeRootHash trie in flip MerkleTrie trie <$> rootHash
 
-type MerkleProof = (Digest SHA256, [Digest SHA256])
+computeHash :: (Show a) => a -> Digest Blake2b_256
+computeHash = hash . pack . show
+
+data MerkleProof k v = MerkleProof
+  { targetKey :: k,
+    targetValue :: v,
+    keyPath :: [(TrieHeight k, k)],
+    siblingHashes :: [Digest Blake2b_256]
+  }
 
 -- | Generate a Merkle proof for a given key in the trie
--- TODO: check the implementation. Right now it's going through the whole trie and picks the hashes of the nodes
--- that match the key. In the branch, when the branch is one of the branches that lead to the root, it computes
--- the hash using the hash of the leaf and the trie on the other side. Is this right?
--- If not, might need to use paramorphism here, as it allows to see the nodes we're operating on alongside the accumulator.
--- That will allow us to properly calculate the hash of the trie from where we came to the current branch.
-proof :: forall k v. (Show k, Show (TrieHeight k), Show v, TrieKey k) => k -> Trie k v -> Maybe MerkleProof
-proof k t = cata go t
+-- Yeah, this variant always constructs whole proof even if the key is not there
+proof :: forall k v. (Show k, Show (TrieHeight k), Show v, TrieKey k) => k -> Trie k v -> Maybe (MerkleProof k v)
+proof k t = let r = para go t in if fst r then snd r else Nothing
   where
-    go :: Algebra (TrieF' k v) (Maybe MerkleProof)
-    go Empty = Nothing
-    -- for every leaf we just compute it's hash
-    go Leaf {..} = Just (computeHash (key, value), [])
-    -- for every branch we look at the prefix
-    go Branch {..}
-      -- if the key doesn't match the prefix, we calculate the hash and store it
-      | not (matchPrefix k prefix (heightToBBit height)) = case (left, right) of
-        (Just (hl, _), Just (hr, _)) -> Just (computeHash (height, prefix, hl, hr), [])
-        (_, _) -> Nothing
-      -- if the key matches the prefix and we're coming from the left branch(zeroBit check),
-      -- then we take the hash from there and append the hash of the current node to the path
-      | zeroBit k (heightToBBit height) = case (left, right) of
-        (Just (hl, p), Just (hr, _)) -> Just (hl, (computeHash (height, prefix, hl, hr)) : p)
-        _ -> Nothing
-      -- the same here, just for the right branch
-      | otherwise = case (left, right) of
-        (Just (hl, _), Just (hr, p)) -> Just (hr, (computeHash (height, prefix, hl, hr)) : p)
-        _ -> Nothing
+    -- Bool to signify if the key was in the structure
+    go :: RAlgebra (TrieF' k v) (Bool, Maybe (MerkleProof k v))
+    go Empty = (False, Nothing)
+    -- for every leaf we just compute merkle proof
+    go Leaf {..} =
+      let targetKey = k
+          targetValue = value
+          keyPath = []
+          siblingHashes = [computeHash (key, value)]
+       in (targetKey == k, Just $ MerkleProof {..})
+    go (Branch h p (_, (bl, Just pl)) (_, (br, Just pr)))
+      -- if one of the proofs containt the key, it means it comes from a path that we're interested in
+      -- it also means, that this node is in the path
+      | pl.targetKey == k =
+          let targetKey = pl.targetKey
+              targetValue = pl.targetValue
+              keyPath = (h, p) : pl.keyPath
+              siblingHashes = pr.siblingHashes <> pl.siblingHashes
+           in (bl, Just $ MerkleProof {..})
+      | pr.targetKey == k =
+          let targetKey = pr.targetKey
+              targetValue = pr.targetValue
+              keyPath = (h, p) : pr.keyPath
+              siblingHashes = pl.siblingHashes <> pr.siblingHashes
+           in (br, Just $ MerkleProof {..})
+      -- otherwise it doesn't matter what key we pick, we're interested only in hash
+      | otherwise =
+          let targetKey = pl.targetKey
+              targetValue = pl.targetValue
+              keyPath = []
+              -- in this case there is always one element in the path list
+              siblingHashes = [computeHash (h, p, head pl.siblingHashes, head pr.siblingHashes)]
+           in (False, Just $ MerkleProof {..})
+    go Branch {} = (False, Nothing)
 
 -- | Validate a Merkle proof against the root hash of the trie
-validate :: Digest SHA256 -> MerkleProof -> Bool
-validate rootHash (leafHash, pf) = foldl go leafHash pf == rootHash
-  where
-    go :: Digest SHA256 -> Digest SHA256 -> Digest SHA256
-    go acc h = computeHash (acc, h)
+validate :: (Show k, Show (TrieHeight k), Show v) => MerkleProof k v -> Digest Blake2b_256 -> Bool
+validate MerkleProof {..} rootHash =
+  rootHash
+    == ( foldr
+           (\((h, p), hs) acc -> computeHash (h, p, hs, acc))
+           (computeHash (targetKey, targetValue))
+           . reverse
+           $ zip keyPath siblingHashes
+       )
diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index 5e4704b..b5f4a2c 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -1,21 +1,50 @@
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
-module Data.Trie (TrieF (..), TrieF', Trie, TrieKey (..), lookup, insertWith, insert, singleton, empty, module X) where
-
--- importing like this to re-export
+module Data.Trie
+  ( TrieF (..),
+    TrieF',
+    Trie,
+    TrieKey (..),
+    prettyPrint,
+    lookup,
+    insertWith,
+    insert,
+    singleton,
+    empty,
+    module X,
+  )
+where
 
 import Data.Internal.RecursionSchemes as X
 import Data.Internal.Trie
--- import Data.Internal.Zipper as X
 import Prelude hiding (lookup)
 
+prettyPrint :: forall k v. (TrieKey k, Show k, Show (TrieHeight k), Show v) => Trie k v -> String
+prettyPrint = cata go
+  where
+    go :: Algebra (TrieF' k v) String
+    go Empty = "()"
+    go Leaf {..} = "(" <> show key <> ", " <> show value <> ")"
+    go Branch {..} =
+      "{"
+        <> show height
+        <> ", "
+        <> show prefix
+        <> "}\n"
+        <> "/"
+        <> replicate 10 ' '
+        <> "\\\n"
+        <> left
+        <> replicate 10 ' '
+        <> right
+
 lookup :: (TrieKey k) => k -> Trie k v -> Maybe v
 lookup _ (In Empty) = Nothing
 lookup k (In Leaf {..}) = if k == key then Just value else Nothing
 lookup k (In Branch {..})
   | not (matchPrefix k prefix (heightToBBit height)) = Nothing
-  | k <= heightToBBit height = lookup k left
+  | zeroBit k (heightToBBit height) = lookup k left
   | otherwise = lookup k right
 
 -- | Inserts new value in the trie resolving possible conflicts using `(v -> v -> v)`.
@@ -32,7 +61,7 @@ insertWith resolve k v t = go t
     go (In Branch {..}) =
       if matchPrefix k prefix $ heightToBBit height
         then
-          if k <= heightToBBit height
+          if zeroBit k $ heightToBBit height
             then In $ Branch height prefix (go left) right
             else In $ Branch height prefix left (go right)
         else join k (In $ Leaf k v) prefix t
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index 2fa77ef..76fdb39 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -62,7 +62,6 @@ library
         Data.Trie
         -- Data.Internal
         Data.Internal.Trie
-        Data.Internal.Zipper
         Data.Internal.RecursionSchemes
         Data.Utils
 
diff --git a/merkle-patricia-trie/test/TrieKeySpec.hs b/merkle-patricia-trie/test/TrieKeySpec.hs
index dd90ffe..aa2af2e 100644
--- a/merkle-patricia-trie/test/TrieKeySpec.hs
+++ b/merkle-patricia-trie/test/TrieKeySpec.hs
@@ -22,30 +22,30 @@ spec :: Spec
 spec = do
   describe "TrieKey typeclass" $ do
     it "mask should discard bits not in the mask" $ do
-      mask @Word256 0b1101 0b0 `shouldBe` 0b0 -- zero mask should return zero
-      mask @Word256 0b1101 0b0100 `shouldBe` 0b1100
-      mask @Word256 0b1101 0b0010 `shouldBe` 0b1100
-      mask @Word256 0b1101 0b0001 `shouldBe` 0b1101
+      mask @Word256 0b1101 0b1000 `shouldBe` 0b0101
+      mask @Word256 0b1101 0b0100 `shouldBe` 0b0001
+      mask @Word256 0b1101 0b0010 `shouldBe` 0b0001
 
     it "commonBranchingBit should find the first bit on which prefixes disagree" $ do
       commonBranchingBit @Word256 0b1100 0b1000 `shouldBe` 0b0100
-      commonBranchingBit @Word256 0b1100 0b1110 `shouldBe` 0b0010
+      commonBranchingBit @Word256 0b1100 0b0100 `shouldBe` 0b1000
 
     it "matchPrefix should mask the key using supplied branching bit and compare to prefix" $ do
-      matchPrefix @Word256 0b1101 0b1100 0b0100 `shouldBe` True
-      matchPrefix @Word256 0b1101 0b1000 0b0100 `shouldBe` False
+      matchPrefix @Word256 0b1101 0b101 0b1000 `shouldBe` True
+      matchPrefix @Word256 0b1101 0b101 0b0100 `shouldBe` False
+      matchPrefix @Word256 0b1101 0b001 0b0100 `shouldBe` True
 
     it "heightToBBit should convert height Word256 to the branching bit" $ do
-      heightToBBit @Word256 0 `shouldBe` setBit 0 255
-      heightToBBit @Word256 1 `shouldBe` setBit 0 254
+      heightToBBit @Word256 1 `shouldBe` 0b00010
+      heightToBBit @Word256 4 `shouldBe` 0b10000
 
     it "bBitToHeight should convert branching bit Word256 to height" $ do
-      bBitToHeight @Word256 (setBit 0 255) `shouldBe` 0
-      bBitToHeight @Word256 (setBit 0 254) `shouldBe` 1
+      bBitToHeight @Word256 0b100 `shouldBe` 2
+      bBitToHeight @Word256 0b010 `shouldBe` 1
 
     it "zeroBit should test whether the desired bit is zero" $ do
-      zeroBit @Word256 0b1101 (heightToBBit @Word256 1) `shouldBe` True
-      zeroBit @Word256 0b1101 (heightToBBit @Word256 2) `shouldBe` False
+      zeroBit @Word256 0b1101 0b0100 `shouldBe` False
+      zeroBit @Word256 0b1101 0b0010 `shouldBe` True
 
     it "QuickCheck property: heightToBBit and bBitToHeight should be inverses" $
       property $

From ddc52a2fda2f8f8cbb1930c46468f006c792371d Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 6 Jan 2025 15:48:46 +0200
Subject: [PATCH 29/31] feat: merkle proof and validation tests

---
 merkle-patricia-trie/lib/Data/MerkleTrie.hs   | 66 +++++++++++--------
 merkle-patricia-trie/lib/Data/Trie.hs         |  6 +-
 merkle-patricia-trie/lib/Data/Utils.hs        |  1 -
 .../merkle-patricia-trie.cabal                |  5 +-
 merkle-patricia-trie/test/MerkleTrieSpec.hs   | 52 +++++++++++++++
 5 files changed, 97 insertions(+), 33 deletions(-)
 delete mode 100644 merkle-patricia-trie/lib/Data/Utils.hs
 create mode 100644 merkle-patricia-trie/test/MerkleTrieSpec.hs

diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index cac8c13..e174088 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -1,6 +1,10 @@
 {-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+-- need this because we need to specify type family in the deriving context
+{-# LANGUAGE StandaloneDeriving #-}
+-- need this to allow type family in the deriving context
+{-# LANGUAGE UndecidableInstances #-}
 
 module Data.MerkleTrie
   ( proof,
@@ -21,20 +25,6 @@ data MerkleTrie k v = MerkleTrie
     trie :: Trie k v
   }
 
-computeRootHash :: forall k v. (Show k, Show v, Show (TrieHeight k)) => Trie k v -> Maybe (Digest Blake2b_256)
-computeRootHash = cata go
-  where
-    go :: Algebra (TrieF' k v) (Maybe (Digest Blake2b_256))
-    go Empty = Nothing
-    go Leaf {..} = Just $ computeHash (key, value)
-    go Branch {..} = Just $ computeHash (height, prefix, left, right)
-
-merkelize :: (Show k, Show v, Show (TrieHeight k)) => Trie k v -> Maybe (MerkleTrie k v)
-merkelize trie = let rootHash = computeRootHash trie in flip MerkleTrie trie <$> rootHash
-
-computeHash :: (Show a) => a -> Digest Blake2b_256
-computeHash = hash . pack . show
-
 data MerkleProof k v = MerkleProof
   { targetKey :: k,
     targetValue :: v,
@@ -42,22 +32,40 @@ data MerkleProof k v = MerkleProof
     siblingHashes :: [Digest Blake2b_256]
   }
 
+deriving instance (Eq k, Eq v, Eq (TrieHeight k)) => Eq (MerkleProof k v)
+
+deriving instance (Show k, Show v, Show (TrieHeight k)) => Show (MerkleProof k v)
+
+computeHash :: (Show a) => a -> Digest Blake2b_256
+computeHash = hash . pack . show
+
+computeRootHash :: forall k v. (Show k, Show v, Show (TrieHeight k)) => Trie k v -> Digest Blake2b_256
+computeRootHash = cata go
+  where
+    go :: Algebra (TrieF' k v) (Digest Blake2b_256)
+    go Empty = computeHash "EmptyTrie"
+    go Leaf {..} = computeHash (key, value)
+    go Branch {..} = computeHash (height, prefix, left, right)
+
+merkelize :: (Show k, Show v, Show (TrieHeight k)) => Trie k v -> MerkleTrie k v
+merkelize trie = let rootHash = computeRootHash trie in MerkleTrie rootHash trie
+
 -- | Generate a Merkle proof for a given key in the trie
 -- Yeah, this variant always constructs whole proof even if the key is not there
 proof :: forall k v. (Show k, Show (TrieHeight k), Show v, TrieKey k) => k -> Trie k v -> Maybe (MerkleProof k v)
-proof k t = let r = para go t in if fst r then snd r else Nothing
+proof k t = let r = cata go t in if fst r then snd r else Nothing
   where
     -- Bool to signify if the key was in the structure
-    go :: RAlgebra (TrieF' k v) (Bool, Maybe (MerkleProof k v))
+    go :: Algebra (TrieF' k v) (Bool, Maybe (MerkleProof k v))
     go Empty = (False, Nothing)
     -- for every leaf we just compute merkle proof
     go Leaf {..} =
-      let targetKey = k
+      let targetKey = key
           targetValue = value
           keyPath = []
-          siblingHashes = [computeHash (key, value)]
-       in (targetKey == k, Just $ MerkleProof {..})
-    go (Branch h p (_, (bl, Just pl)) (_, (br, Just pr)))
+          siblingHashes = [computeHash (key, value) | key /= k]
+       in (key == k, Just $ MerkleProof {..})
+    go (Branch h p (bl, Just pl) (br, Just pr))
       -- if one of the proofs containt the key, it means it comes from a path that we're interested in
       -- it also means, that this node is in the path
       | pl.targetKey == k =
@@ -70,7 +78,7 @@ proof k t = let r = para go t in if fst r then snd r else Nothing
           let targetKey = pr.targetKey
               targetValue = pr.targetValue
               keyPath = (h, p) : pr.keyPath
-              siblingHashes = pl.siblingHashes <> pr.siblingHashes
+              siblingHashes = pr.siblingHashes <> pl.siblingHashes
            in (br, Just $ MerkleProof {..})
       -- otherwise it doesn't matter what key we pick, we're interested only in hash
       | otherwise =
@@ -83,12 +91,14 @@ proof k t = let r = para go t in if fst r then snd r else Nothing
     go Branch {} = (False, Nothing)
 
 -- | Validate a Merkle proof against the root hash of the trie
-validate :: (Show k, Show (TrieHeight k), Show v) => MerkleProof k v -> Digest Blake2b_256 -> Bool
+validate :: (Show k, Show (TrieHeight k), Show v, TrieKey k) => MerkleProof k v -> Digest Blake2b_256 -> Bool
 validate MerkleProof {..} rootHash =
   rootHash
-    == ( foldr
-           (\((h, p), hs) acc -> computeHash (h, p, hs, acc))
-           (computeHash (targetKey, targetValue))
-           . reverse
-           $ zip keyPath siblingHashes
-       )
+    == foldr
+      ( \((h, p), hs) acc ->
+          if zeroBit targetKey (heightToBBit h)
+            then computeHash (h, p, acc, hs)
+            else computeHash (h, p, hs, acc)
+      )
+      (computeHash (targetKey, targetValue))
+      (reverse $ zip keyPath siblingHashes)
diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index b5f4a2c..1dc0e93 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -27,13 +27,13 @@ prettyPrint = cata go
     go Empty = "()"
     go Leaf {..} = "(" <> show key <> ", " <> show value <> ")"
     go Branch {..} =
-      "{"
+      "{h:"
         <> show height
-        <> ", "
+        <> ", p:"
         <> show prefix
         <> "}\n"
         <> "/"
-        <> replicate 10 ' '
+        <> replicate 20 ' '
         <> "\\\n"
         <> left
         <> replicate 10 ' '
diff --git a/merkle-patricia-trie/lib/Data/Utils.hs b/merkle-patricia-trie/lib/Data/Utils.hs
deleted file mode 100644
index 19e622d..0000000
--- a/merkle-patricia-trie/lib/Data/Utils.hs
+++ /dev/null
@@ -1 +0,0 @@
-module Data.Utils where
diff --git a/merkle-patricia-trie/merkle-patricia-trie.cabal b/merkle-patricia-trie/merkle-patricia-trie.cabal
index 76fdb39..2da51c0 100644
--- a/merkle-patricia-trie/merkle-patricia-trie.cabal
+++ b/merkle-patricia-trie/merkle-patricia-trie.cabal
@@ -63,7 +63,6 @@ library
         -- Data.Internal
         Data.Internal.Trie
         Data.Internal.RecursionSchemes
-        Data.Utils
 
     -- Modules included in this library but not exported.
     -- other-modules:
@@ -96,6 +95,10 @@ test-suite test
     type: exitcode-stdio-1.0
     hs-source-dirs: test
     main-is: Spec.hs
+    other-modules:
+        MerkleTrieSpec
+        TrieSpec
+        TrieKeySpec
     build-depends:
         base,
         merkle-patricia-trie,
diff --git a/merkle-patricia-trie/test/MerkleTrieSpec.hs b/merkle-patricia-trie/test/MerkleTrieSpec.hs
new file mode 100644
index 0000000..463f195
--- /dev/null
+++ b/merkle-patricia-trie/test/MerkleTrieSpec.hs
@@ -0,0 +1,52 @@
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module MerkleTrieSpec (spec) where
+
+import Data.MerkleTrie
+import Data.Trie qualified as Trie
+import Data.WideWord (Word256)
+import Data.Word (Word8)
+import Test.Hspec
+import Test.QuickCheck
+
+instance Trie.TrieKey Word256 where
+  type TrieHeight Word256 = Word8
+
+spec :: Spec
+spec = do
+  describe "MerkleTrie" $ do
+    it "should generate a proof, validate it, and compute the root hash correctly" $ do
+      let t = Trie.insert @Word256 @_ 1 "value1" $ Trie.insert 2 "value2" Trie.empty
+          merkleTrie = merkelize t
+          proof1 = proof 1 t
+          proof2 = proof 2 t
+      putStrLn $ Trie.prettyPrint t
+      print proof1
+      print proof2
+      case (proof1, proof2) of
+        (Just p1, Just p2) -> do
+          validate p1 (rootHash merkleTrie) `shouldBe` True
+          validate p2 (rootHash merkleTrie) `shouldBe` True
+          computeRootHash t `shouldBe` rootHash merkleTrie
+        _ -> expectationFailure "Failed to generate proofs"
+
+    it "should return Nothing for a proof of a key not in the trie" $ do
+      let t = Trie.insert @Word256 @_ 1 "value1" $ Trie.insert 2 "value2" Trie.empty
+      proof 3 t `shouldBe` Nothing
+
+    it "should fail to validate an incorrect proof" $ do
+      let t = Trie.insert @Word256 @_ 1 "value1" $ Trie.insert 2 "value2" Trie.empty
+          merkleTrie = merkelize t
+          proof1 = proof 1 t
+          proof2 = proof 2 t
+      case (proof1, proof2) of
+        (Just p1, Just p2) -> do
+          validate p1 (rootHash merkleTrie) `shouldBe` True
+          validate p2 (rootHash merkleTrie) `shouldBe` True
+          -- Modify proof1 to make it invalid
+          let invalidProof = p1 {targetValue = "invalidValue"}
+          validate invalidProof (rootHash merkleTrie) `shouldBe` False
+        _ -> expectationFailure "Failed to generate proofs"

From d99bfd72eee52ea86a4ed6ec794001da8c37dc2f Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Mon, 6 Jan 2025 16:34:05 +0200
Subject: [PATCH 30/31] fix: small refactoring; more tests

---
 merkle-patricia-trie/lib/Data/Trie.hs       | 20 --------------------
 merkle-patricia-trie/test/MerkleTrieSpec.hs |  4 ----
 merkle-patricia-trie/test/TrieSpec.hs       |  7 ++++++-
 3 files changed, 6 insertions(+), 25 deletions(-)

diff --git a/merkle-patricia-trie/lib/Data/Trie.hs b/merkle-patricia-trie/lib/Data/Trie.hs
index 1dc0e93..6a14634 100644
--- a/merkle-patricia-trie/lib/Data/Trie.hs
+++ b/merkle-patricia-trie/lib/Data/Trie.hs
@@ -6,7 +6,6 @@ module Data.Trie
     TrieF',
     Trie,
     TrieKey (..),
-    prettyPrint,
     lookup,
     insertWith,
     insert,
@@ -20,25 +19,6 @@ import Data.Internal.RecursionSchemes as X
 import Data.Internal.Trie
 import Prelude hiding (lookup)
 
-prettyPrint :: forall k v. (TrieKey k, Show k, Show (TrieHeight k), Show v) => Trie k v -> String
-prettyPrint = cata go
-  where
-    go :: Algebra (TrieF' k v) String
-    go Empty = "()"
-    go Leaf {..} = "(" <> show key <> ", " <> show value <> ")"
-    go Branch {..} =
-      "{h:"
-        <> show height
-        <> ", p:"
-        <> show prefix
-        <> "}\n"
-        <> "/"
-        <> replicate 20 ' '
-        <> "\\\n"
-        <> left
-        <> replicate 10 ' '
-        <> right
-
 lookup :: (TrieKey k) => k -> Trie k v -> Maybe v
 lookup _ (In Empty) = Nothing
 lookup k (In Leaf {..}) = if k == key then Just value else Nothing
diff --git a/merkle-patricia-trie/test/MerkleTrieSpec.hs b/merkle-patricia-trie/test/MerkleTrieSpec.hs
index 463f195..76a70c4 100644
--- a/merkle-patricia-trie/test/MerkleTrieSpec.hs
+++ b/merkle-patricia-trie/test/MerkleTrieSpec.hs
@@ -23,14 +23,10 @@ spec = do
           merkleTrie = merkelize t
           proof1 = proof 1 t
           proof2 = proof 2 t
-      putStrLn $ Trie.prettyPrint t
-      print proof1
-      print proof2
       case (proof1, proof2) of
         (Just p1, Just p2) -> do
           validate p1 (rootHash merkleTrie) `shouldBe` True
           validate p2 (rootHash merkleTrie) `shouldBe` True
-          computeRootHash t `shouldBe` rootHash merkleTrie
         _ -> expectationFailure "Failed to generate proofs"
 
     it "should return Nothing for a proof of a key not in the trie" $ do
diff --git a/merkle-patricia-trie/test/TrieSpec.hs b/merkle-patricia-trie/test/TrieSpec.hs
index d7517c7..e2bcfe0 100644
--- a/merkle-patricia-trie/test/TrieSpec.hs
+++ b/merkle-patricia-trie/test/TrieSpec.hs
@@ -1,10 +1,11 @@
+{-# LANGUAGE ImportQualifiedPost #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 
 module TrieSpec (spec) where
 
-import qualified Data.Trie as Trie
+import Data.Trie qualified as Trie
 import Data.WideWord (Word256)
 import Data.Word (Word, Word8)
 import Test.Hspec
@@ -59,3 +60,7 @@ spec = do
     it "insertWith should handle conflicts correctly" $ do
       let t = Trie.insertWith @Word256 @_ (++) 1 "value1" $ Trie.insertWith (++) 1 "value2" Trie.empty
       Trie.lookup 1 t `shouldBe` Just "value1value2"
+
+    it "QuickCheck property: insert/lookup roundtrip for randomly generated tries" $ property $ \(keys :: [Word]) -> do
+      let t = foldr (\k tr -> Trie.insert @Word256 @_ (fromIntegral k) (show k) tr) Trie.empty keys
+      all (\k -> Trie.lookup (fromIntegral k) t == Just (show k)) keys `shouldBe` True

From d255bafe7478fa371984d5a3984a14dbdb5652c1 Mon Sep 17 00:00:00 2001
From: Yaroslav Kozhevnikov <yroslav541@gmail.com>
Date: Tue, 7 Jan 2025 13:18:40 +0200
Subject: [PATCH 31/31] chore: todos

---
 merkle-patricia-trie/lib/Data/MerkleTrie.hs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/merkle-patricia-trie/lib/Data/MerkleTrie.hs b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
index e174088..5ab1ad1 100644
--- a/merkle-patricia-trie/lib/Data/MerkleTrie.hs
+++ b/merkle-patricia-trie/lib/Data/MerkleTrie.hs
@@ -36,6 +36,7 @@ deriving instance (Eq k, Eq v, Eq (TrieHeight k)) => Eq (MerkleProof k v)
 
 deriving instance (Show k, Show v, Show (TrieHeight k)) => Show (MerkleProof k v)
 
+-- TODO: replace the 'show' with binary serialization
 computeHash :: (Show a) => a -> Digest Blake2b_256
 computeHash = hash . pack . show