index hashtable with normalized material keys

niklasf · Dec 30, 2024 · e05ec10 · e05ec10
1 parent 188d01a
commit e05ec10
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 40 deletions.
diff --git a/src/errors.rs b/src/errors.rs
@@ -1,6 +1,9 @@
 use std::{backtrace::Backtrace, error::Error, fmt, io};
 
-use crate::{material::Material, types::Metric};
+use crate::{
+    material::{Material, NormalizedMaterial},
+    types::Metric,
+};
 
 pub type SyzygyResult<T> = Result<T, SyzygyError>;
 
@@ -102,14 +105,14 @@ impl Error for ProbeError {
 }
 
 pub trait ProbeResultExt<T> {
-    fn ctx(self, metric: Metric, material: Material) -> SyzygyResult<T>;
+    fn ctx(self, metric: Metric, material: &NormalizedMaterial) -> SyzygyResult<T>;
 }
 
 impl<T> ProbeResultExt<T> for ProbeResult<T> {
-    fn ctx(self, metric: Metric, material: Material) -> SyzygyResult<T> {
+    fn ctx(self, metric: Metric, material: &NormalizedMaterial) -> SyzygyResult<T> {
         self.map_err(|error| SyzygyError::ProbeFailed {
             metric,
-            material: material.into_normalized(),
+            material: material.inner().clone(),
             error: Box::new(error),
         })
     }

diff --git a/src/material.rs b/src/material.rs
@@ -158,10 +158,10 @@ impl Material {
         }
     }
 
-    pub(crate) fn into_normalized(self) -> Material {
-        Material {
-            by_color: self.by_color.into_normalized(),
-        }
+    pub(crate) fn to_normalized(&self) -> NormalizedMaterial {
+        NormalizedMaterial(Material {
+            by_color: self.by_color.clone().into_normalized(),
+        })
     }
 }
 
@@ -170,3 +170,12 @@ impl fmt::Display for Material {
         write!(f, "{}v{}", self.by_color.white, self.by_color.black)
     }
 }
+
+#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+pub(crate) struct NormalizedMaterial(Material);
+
+impl NormalizedMaterial {
+    pub fn inner(&self) -> &Material {
+        &self.0
+    }
+}
diff --git a/src/tablebase.rs b/src/tablebase.rs
@@ -16,7 +16,7 @@ use crate::{
     errors::{ProbeResultExt as _, SyzygyError, SyzygyResult},
     filesystem,
     filesystem::Filesystem,
-    material::Material,
+    material::{Material, NormalizedMaterial},
     table::{DtzTable, WdlTable},
     types::{DecisiveWdl, Dtz, MaybeRounded, Metric, Syzygy, Wdl},
     AmbiguousWdl,
@@ -37,8 +37,8 @@ enum ProbeState {
 /// A collection of tables.
 pub struct Tablebase<S: Position + Clone + Syzygy> {
     filesystem: Arc<dyn Filesystem>,
-    wdl: FxHashMap<Material, (PathBuf, OnceCell<WdlTable<S>>)>,
-    dtz: FxHashMap<Material, (PathBuf, OnceCell<DtzTable<S>>)>,
+    wdl: FxHashMap<NormalizedMaterial, (PathBuf, OnceCell<WdlTable<S>>)>,
+    dtz: FxHashMap<NormalizedMaterial, (PathBuf, OnceCell<DtzTable<S>>)>,
     max_pieces: usize,
 }
 
@@ -204,47 +204,39 @@ impl<S: Position + Clone + Syzygy> Tablebase<S> {
         }
 
         // Add path.
+        let material = material.to_normalized();
+        let path = path.to_path_buf();
         if is_tbw {
-            self.wdl
-                .insert(material, (path.to_path_buf(), OnceCell::new()));
+            self.wdl.insert(material, (path, OnceCell::new()));
         } else {
-            self.dtz
-                .insert(material, (path.to_path_buf(), OnceCell::new()));
+            self.dtz.insert(material, (path, OnceCell::new()));
         }
         self.max_pieces = max(self.max_pieces, pieces);
         Ok(())
     }
 
-    fn wdl_table(&self, key: &Material) -> SyzygyResult<&WdlTable<S>> {
-        if let Some((path, table)) = self
-            .wdl
-            .get(key)
-            .or_else(|| self.wdl.get(&key.clone().into_swapped()))
-        {
+    fn wdl_table(&self, material: &NormalizedMaterial) -> SyzygyResult<&WdlTable<S>> {
+        if let Some((path, table)) = self.wdl.get(material) {
             table
-                .get_or_try_init(|| WdlTable::new(self.filesystem.open(path)?, key))
-                .ctx(Metric::Wdl, key.clone())
+                .get_or_try_init(|| WdlTable::new(self.filesystem.open(path)?, material.inner()))
+                .ctx(Metric::Wdl, material)
         } else {
             Err(SyzygyError::MissingTable {
                 metric: Metric::Wdl,
-                material: key.clone().into_normalized(),
+                material: material.inner().clone(),
             })
         }
     }
 
-    fn dtz_table(&self, key: &Material) -> SyzygyResult<&DtzTable<S>> {
-        if let Some((path, table)) = self
-            .dtz
-            .get(key)
-            .or_else(|| self.dtz.get(&key.clone().into_swapped()))
-        {
+    fn dtz_table(&self, material: &NormalizedMaterial) -> SyzygyResult<&DtzTable<S>> {
+        if let Some((path, table)) = self.dtz.get(material) {
             table
-                .get_or_try_init(|| DtzTable::new(self.filesystem.open(path)?, key))
-                .ctx(Metric::Dtz, key.clone())
+                .get_or_try_init(|| DtzTable::new(self.filesystem.open(path)?, material.inner()))
+                .ctx(Metric::Dtz, material)
         } else {
             Err(SyzygyError::MissingTable {
                 metric: Metric::Dtz,
-                material: key.clone().into_normalized(),
+                material: material.inner().clone(),
             })
         }
     }
@@ -643,9 +635,9 @@ impl<S: Position + Clone + Syzygy> Tablebase<S> {
         }
 
         // Get raw WDL value from the appropriate table.
-        let key = Material::from_board(pos.board());
-        self.wdl_table(&key)
-            .and_then(|table| table.probe_wdl(pos).ctx(Metric::Wdl, key))
+        let material = Material::from_board(pos.board()).to_normalized();
+        self.wdl_table(&material)
+            .and_then(|table| table.probe_wdl(pos).ctx(Metric::Wdl, &material))
     }
 
     fn probe_dtz_table(
@@ -654,9 +646,9 @@ impl<S: Position + Clone + Syzygy> Tablebase<S> {
         wdl: DecisiveWdl,
     ) -> SyzygyResult<Option<MaybeRounded<u32>>> {
         // Get raw DTZ value from the appropriate table.
-        let key = Material::from_board(pos.board());
-        self.dtz_table(&key)
-            .and_then(|table| table.probe_dtz(pos, wdl).ctx(Metric::Dtz, key))
+        let material = Material::from_board(pos.board()).to_normalized();
+        self.dtz_table(&material)
+            .and_then(|table| table.probe_dtz(pos, wdl).ctx(Metric::Dtz, &material))
     }
 }
 
@@ -739,7 +731,10 @@ impl<'a, S: Position + Clone + Syzygy + 'a> WdlEntry<'a, S> {
             }
         }
 
-        (|| Ok(u!(best)))().ctx(Metric::Dtz, Material::from_board(self.pos.board()))
+        (|| Ok(u!(best)))().ctx(
+            Metric::Dtz,
+            &Material::from_board(self.pos.board()).to_normalized(),
+        )
     }
 }