feat: remove DistanceFunction (#1838)

Author: morph-dev

crates/ethportal-api/src/types/distance.rs

@@ -82,6 +82,7 @@ pub trait Metric {
 }
 
 /// The XOR metric defined in the Kademlia paper.
+#[derive(Debug)]
 pub struct XorMetric;
 
 impl Metric for XorMetric {

crates/portalnet/src/overlay/service/manager.rs

@@ -977,7 +977,7 @@ mod tests {
     use tokio::{sync::mpsc::unbounded_channel, time::timeout};
     use tokio_test::{assert_pending, assert_ready, task};
     use trin_metrics::portalnet::PORTALNET_METRICS;
-    use trin_storage::{DistanceFunction, MemoryContentStore};
+    use trin_storage::MemoryContentStore;
     use trin_validation::validator::MockValidator;
 
     use super::*;
@@ -1028,7 +1028,7 @@ mod tests {
         let utp_controller = Arc::new(utp_controller);
 
         let node_id = discovery.local_enr().node_id();
-        let store = MemoryContentStore::new(node_id, DistanceFunction::Xor);
+        let store = MemoryContentStore::new(node_id);
         let store = Arc::new(Mutex::new(store));
 
         let overlay_config = OverlayConfig::default();

crates/portalnet/tests/overlay.rs

@@ -30,7 +30,7 @@ use tokio::{
     sync::{mpsc, mpsc::unbounded_channel},
     time::{self, Duration},
 };
-use trin_storage::{ContentStore, DistanceFunction, MemoryContentStore};
+use trin_storage::{ContentStore, MemoryContentStore};
 use trin_validation::validator::MockValidator;
 use utp_rs::socket::UtpSocket;
 
@@ -47,7 +47,7 @@ async fn init_overlay(
     let overlay_config = OverlayConfig::default();
 
     let node_id = discovery.local_enr().node_id();
-    let store = MemoryContentStore::new(node_id, DistanceFunction::Xor);
+    let store = MemoryContentStore::new(node_id);
     let store = Arc::new(Mutex::new(store));
 
     let (_utp_talk_req_tx, utp_talk_req_rx) = unbounded_channel();

crates/storage/src/config.rs

@@ -5,7 +5,7 @@ use ethportal_api::types::{distance::Distance, network::Subnetwork};
 use r2d2::Pool;
 use r2d2_sqlite::SqliteConnectionManager;
 
-use crate::{error::ContentStoreError, utils::setup_sql, DistanceFunction};
+use crate::{error::ContentStoreError, utils::setup_sql};
 
 const BYTES_IN_MB_U64: u64 = 1000 * 1000;
 
@@ -103,7 +103,6 @@ impl PortalStorageConfigFactory {
             storage_capacity_bytes: capacity_bytes,
             node_id: self.node_id,
             node_data_dir: self.node_data_dir.clone(),
-            distance_fn: DistanceFunction::Xor,
             sql_connection_pool: self.sql_connection_pool.clone(),
             max_radius,
         })
@@ -124,7 +123,6 @@ pub struct PortalStorageConfig {
     pub storage_capacity_bytes: u64,
     pub node_id: NodeId,
     pub node_data_dir: PathBuf,
-    pub distance_fn: DistanceFunction,
     pub sql_connection_pool: Pool<SqliteConnectionManager>,
     pub max_radius: Distance,
 }

crates/storage/src/lib.rs

@@ -5,7 +5,7 @@ pub mod test_utils;
 pub mod utils;
 pub mod versioned;
 
-use std::{ops::Deref, str::FromStr};
+use std::{marker::PhantomData, ops::Deref, str::FromStr};
 
 use alloy::primitives::{Bytes, B256};
 pub use config::{PortalStorageConfig, PortalStorageConfigFactory};
@@ -22,21 +22,6 @@ use rusqlite::types::{FromSql, FromSqlError, ValueRef};
 
 pub const DATABASE_NAME: &str = "trin.sqlite";
 
-// TODO: Replace enum with generic type parameter. This will require that we have a way to
-// associate a "find farthest" query with the generic Metric.
-#[derive(Copy, Clone, Debug)]
-pub enum DistanceFunction {
-    Xor,
-}
-
-impl DistanceFunction {
-    pub fn distance(&self, node_id: &NodeId, other: &[u8; 32]) -> Distance {
-        match self {
-            DistanceFunction::Xor => XorMetric::distance(&node_id.raw(), other),
-        }
-    }
-}
-
 /// An enum which tells us if we should store or not store content, and if not why for better
 /// errors.
 #[derive(Debug, PartialEq)]
@@ -76,25 +61,25 @@ pub trait ContentStore {
 }
 
 /// An in-memory `ContentStore`.
-pub struct MemoryContentStore {
+pub struct MemoryContentStore<TMetric: Metric = XorMetric> {
     /// The content store.
     store: std::collections::HashMap<Vec<u8>, RawContentValue>,
     /// The `NodeId` of the local node.
     node_id: NodeId,
-    /// The distance function used by the store to compute distances.
-    distance_fn: DistanceFunction,
     /// The radius of the store.
     radius: Distance,
+    /// Phantom metric
+    _metric: PhantomData<TMetric>,
 }
 
-impl MemoryContentStore {
+impl<TMetric: Metric> MemoryContentStore<TMetric> {
     /// Constructs a new `MemoryPortalContentStore`.
-    pub fn new(node_id: NodeId, distance_fn: DistanceFunction) -> Self {
+    pub fn new(node_id: NodeId) -> Self {
         Self {
             store: std::collections::HashMap::new(),
             node_id,
-            distance_fn,
             radius: Distance::MAX,
+            _metric: PhantomData,
         }
     }
 
@@ -105,7 +90,7 @@ impl MemoryContentStore {
 
     /// Returns the distance to `key` from the local `NodeId` according to the distance function.
     fn distance_to_key<K: OverlayContentKey>(&self, key: &K) -> Distance {
-        self.distance_fn.distance(&self.node_id, &key.content_id())
+        TMetric::distance(&self.node_id.raw(), &key.content_id())
     }
 
     /// Returns `true` if the content store contains data for `key`.
@@ -115,7 +100,7 @@ impl MemoryContentStore {
     }
 }
 
-impl ContentStore for MemoryContentStore {
+impl<TMetric: Metric> ContentStore for MemoryContentStore<TMetric> {
     type Key = IdentityContentKey;
 
     fn get(&self, key: &Self::Key) -> Result<Option<RawContentValue>, ContentStoreError> {
@@ -208,14 +193,14 @@ pub struct DataSize {
 #[allow(clippy::unwrap_used)]
 pub mod test {
     use alloy::primitives::{bytes, B512};
-    use ethportal_api::IdentityContentKey;
+    use ethportal_api::{types::distance::XorMetric, IdentityContentKey};
 
     use super::*;
 
     #[test]
     fn memory_store_contains_key() {
         let node_id = NodeId::random();
-        let mut store = MemoryContentStore::new(node_id, DistanceFunction::Xor);
+        let mut store = MemoryContentStore::<XorMetric>::new(node_id);
 
         let val = vec![0xef];
 
@@ -231,7 +216,7 @@ pub mod test {
     #[test]
     fn memory_store_get() {
         let node_id = NodeId::random();
-        let mut store = MemoryContentStore::new(node_id, DistanceFunction::Xor);
+        let mut store = MemoryContentStore::<XorMetric>::new(node_id);
 
         let val = bytes!("ef");
 
@@ -247,7 +232,7 @@ pub mod test {
     #[test]
     fn memory_store_put() {
         let node_id = NodeId::random();
-        let mut store = MemoryContentStore::new(node_id, DistanceFunction::Xor);
+        let mut store = MemoryContentStore::<XorMetric>::new(node_id);
 
         let val = bytes!("ef");
 
@@ -260,7 +245,7 @@ pub mod test {
     #[test]
     fn memory_store_is_within_radius_and_unavailable() {
         let node_id = NodeId::random();
-        let mut store = MemoryContentStore::new(node_id, DistanceFunction::Xor);
+        let mut store = MemoryContentStore::<XorMetric>::new(node_id);
 
         let val = bytes!("ef");
 

crates/storage/src/versioned/id_indexed_v1/config.rs

@@ -6,7 +6,7 @@ use r2d2::Pool;
 use r2d2_sqlite::SqliteConnectionManager;
 
 use super::pruning_strategy::PruningConfig;
-use crate::{versioned::ContentType, DistanceFunction, PortalStorageConfig};
+use crate::{versioned::ContentType, PortalStorageConfig};
 
 /// The config for the IdIndexedV1Store
 #[derive(Clone, Debug)]
@@ -17,7 +17,6 @@ pub struct IdIndexedV1StoreConfig {
     pub node_data_dir: PathBuf,
     pub storage_capacity_bytes: u64,
     pub sql_connection_pool: Pool<SqliteConnectionManager>,
-    pub distance_fn: DistanceFunction,
     pub pruning_config: PruningConfig,
     pub max_radius: Distance,
 }
@@ -35,7 +34,6 @@ impl IdIndexedV1StoreConfig {
             node_data_dir: config.node_data_dir,
             storage_capacity_bytes: config.storage_capacity_bytes,
             sql_connection_pool: config.sql_connection_pool,
-            distance_fn: config.distance_fn,
             // consider making this a parameter if we start using non-default value
             pruning_config: PruningConfig::default(),
             max_radius: config.max_radius,

crates/storage/src/versioned/id_indexed_v1/pruning_strategy.rs

@@ -203,7 +203,7 @@ mod tests {
     use rstest::rstest;
 
     use super::*;
-    use crate::{versioned::ContentType, DistanceFunction};
+    use crate::versioned::ContentType;
 
     const DEFAULT_STORAGE_CAPACITY_BYTES: u64 = 1_000_000;
     const EXTRA_DISK_USAGE_PER_ENTRY: u64 = 100;
@@ -220,7 +220,6 @@ mod tests {
             node_data_dir: PathBuf::default(),
             storage_capacity_bytes,
             sql_connection_pool: Pool::new(SqliteConnectionManager::memory()).unwrap(),
-            distance_fn: DistanceFunction::Xor,
             pruning_config: PruningConfig::default(),
             max_radius: Distance::MAX,
         };