Run rules after deserialization

core-relations/Cargo.toml

@@ -39,6 +39,7 @@ typetag = { workspace = true }
 serde_json = { workspace = true }
 ordered-float = { workspace = true }
 num-bigint = { workspace = true }
+flexbuffers = { workspace = true }
 
 
 [dev-dependencies]

core-relations/src/base_values/mod.rs

@@ -8,8 +8,10 @@ use std::{
 
 use hashbrown::HashMap;
 use num::{rational::Ratio, BigInt, Rational64};
+use once_cell::sync::Lazy;
 use serde::{de, ser::SerializeStruct, Deserialize, Serialize};
 use serde_json::json;
+use std::sync::Mutex;
 
 use crate::numeric_id::{define_id, DenseIdMap, NumericId};
 
@@ -200,6 +202,12 @@ trait DynamicInternTable: Any + dyn_clone::DynClone + Send + Sync {
 // Implements `Clone` for `Box<dyn DynamicInternTable>`.
 dyn_clone::clone_trait_object!(DynamicInternTable);
 
+type DynTableDeserializer =
+    fn(serde_json::Value) -> Result<Box<dyn DynamicInternTable>, Box<dyn std::error::Error>>;
+
+static EXTERNAL_BASE_VALUE_DESERIALIZERS: Lazy<Mutex<HashMap<String, DynTableDeserializer>>> =
+    Lazy::new(|| Mutex::new(HashMap::new()));
+
 #[derive(Serialize, Deserialize)]
 struct BaseInternTableErased {
     table: serde_json::Value,
@@ -219,6 +227,20 @@ impl<P: BaseValue> Default for BaseInternTable<P> {
     }
 }
 
+fn deserialize_table_for<P: BaseValue>(
+    table: serde_json::Value,
+) -> Result<Box<dyn DynamicInternTable>, Box<dyn std::error::Error>> {
+    let table: InternTable<P, Value> = serde_json::from_value(table)?;
+    Ok(Box::new(BaseInternTable { table }))
+}
+
+pub fn register_deserializable_base_type<P: BaseValue>() {
+    EXTERNAL_BASE_VALUE_DESERIALIZERS
+        .lock()
+        .unwrap()
+        .insert(P::type_id_string(), deserialize_table_for::<P>);
+}
+
 impl<P> DynamicInternTable for BaseInternTable<P>
 where
     P: BaseValue + Serialize + 'static,
@@ -415,6 +437,14 @@ fn deserialize_dyn(
     value: serde_json::Value,
 ) -> Result<Box<dyn DynamicInternTable>, Box<dyn std::error::Error>> {
     let erased: BaseInternTableErased = serde_json::from_value(value)?;
+    if let Some(deserialize) = EXTERNAL_BASE_VALUE_DESERIALIZERS
+        .lock()
+        .unwrap()
+        .get(&erased.base_value_type)
+        .copied()
+    {
+        return deserialize(erased.table);
+    }
 
     match erased.base_value_type.as_str() {
         "Unit" => {

core-relations/src/common.rs

@@ -7,9 +7,9 @@ use std::{
 
 use crate::numeric_id::{define_id, DenseIdMap, IdVec, NumericId};
 use egglog_concurrency::ConcurrentVec;
+use flexbuffers::{FlexbufferSerializer, Reader};
 use rustc_hash::FxHasher;
 use serde::{ser::SerializeStruct, Deserialize, Deserializer, Serialize};
-use serde_json::{from_value, to_value};
 
 use crate::{pool::Clear, Subset, TableId, TableVersion, WrappedTable};
 
@@ -43,24 +43,43 @@ where
     where
         S: serde::Serializer,
     {
+        fn to_flexbuffer<T: Serialize>(value: &T) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
+            let mut serializer = FlexbufferSerializer::new();
+            value.serialize(&mut serializer)?;
+            Ok(serializer.take_buffer())
+        }
+
+        let vals = self
+            .vals
+            .read()
+            .iter()
+            .map(|value| to_flexbuffer(value).map_err(serde::ser::Error::custom))
+            .collect::<Result<Vec<_>, _>>()?;
+
         let mut s = serializer.serialize_struct("InternTable", 3)?;
-        s.serialize_field("vals", &self.vals)?;
+        s.serialize_field("vals", &vals)?;
 
-        let mut serializable_data: Vec<Vec<(serde_json::Value, serde_json::Value)>> = Vec::new();
+        #[derive(Serialize)]
+        struct SerializedEntry {
+            key: Vec<u8>,
+            value: Vec<u8>,
+        }
+
+        let mut serializable_data: Vec<Vec<SerializedEntry>> = Vec::new();
         for shard in &self.data {
             let shard = shard
                 .lock()
                 .map_err(|_| serde::ser::Error::custom("mutex poisoned"))?;
             let mut inner = Vec::new();
             for (k, v) in shard.iter() {
-                inner.push((
-                    to_value(k).map_err(serde::ser::Error::custom)?,
-                    to_value(v).map_err(serde::ser::Error::custom)?,
-                ));
+                inner.push(SerializedEntry {
+                    key: to_flexbuffer(k).map_err(serde::ser::Error::custom)?,
+                    value: to_flexbuffer(v).map_err(serde::ser::Error::custom)?,
+                });
             }
 
-            // Sort by the serialized key to guarantee deterministic order
-            inner.sort_by(|(k1, _), (k2, _)| k1.to_string().cmp(&k2.to_string()));
+            // Sort by the encoded key bytes to guarantee deterministic order.
+            inner.sort_by(|e1, e2| e1.key.cmp(&e2.key));
 
             serializable_data.push(inner)
         }
@@ -80,33 +99,52 @@ where
     where
         D: Deserializer<'de>,
     {
-        // Helper struct matching the serialized shape
+        fn from_flexbuffer<T: for<'a> Deserialize<'a>>(
+            bytes: &[u8],
+        ) -> Result<T, Box<dyn std::error::Error>> {
+            let reader = Reader::get_root(bytes)?;
+            Ok(T::deserialize(reader)?)
+        }
+
         #[derive(Deserialize)]
-        struct Partial<K> {
-            vals: Arc<ConcurrentVec<K>>,
-            data: Vec<Vec<(serde_json::Value, serde_json::Value)>>,
+        struct SerializedEntry {
+            key: Vec<u8>,
+            value: Vec<u8>,
+        }
+
+        #[derive(Deserialize)]
+        struct Partial {
+            vals: Vec<Vec<u8>>,
+            data: Vec<Vec<SerializedEntry>>,
             shards_log2: u32,
         }
 
         let helper = Partial::deserialize(deserializer)?;
+        let vals = Arc::new(ConcurrentVec::with_capacity(helper.vals.len().max(128)));
+        for bytes in helper.vals {
+            let value = from_flexbuffer(&bytes).unwrap();
+            vals.push(value);
+        }
 
-        // Convert each shard from Vec<(Value, Value)> back into HashMap<K, V>
         let data: Vec<Arc<Mutex<hashbrown::HashMap<K, V>>>> = helper
             .data
             .into_iter()
             .map(|shard_vec| {
                 let mut map = hashbrown::HashMap::new();
-                for (k_val, v_val) in shard_vec {
-                    let k: K = from_value(k_val).unwrap();
-                    let v: V = from_value(v_val).unwrap();
+                for entry in shard_vec {
+                    let SerializedEntry { key, value } = entry;
+                    let (k, v) = (
+                        from_flexbuffer(&key).unwrap(),
+                        from_flexbuffer(&value).unwrap(),
+                    );
                     map.insert(k, v);
                 }
                 Arc::new(Mutex::new(map))
             })
             .collect();
 
         Ok(InternTable {
-            vals: helper.vals,
+            vals,
             data,
             shards_log2: helper.shards_log2,
         })
@@ -294,6 +332,13 @@ pub(crate) struct SubsetTracker {
 }
 
 impl SubsetTracker {
+    /// Deserialized tables may have different row/version state than when this
+    /// tracker was recorded, so cached "last rebuilt at" entries cannot be
+    /// trusted across serialization boundaries.
+    pub(crate) fn restore_deserialized_runtime(&mut self) {
+        *self = Default::default();
+    }
+
     /// Hand back the subset of the table needed to be scanned in order to see all updates since
     /// the last call to this method.
     ///

core-relations/src/containers/mod.rs

@@ -65,6 +65,10 @@ impl ContainerValues {
         Default::default()
     }
 
+    pub fn restore_deserialized_runtime(&mut self) {
+        self.subset_tracker.restore_deserialized_runtime();
+    }
+
     fn get<C: ContainerValue>(&self) -> Option<&ContainerEnv<C>> {
         let id = self
             .container_ids

core-relations/src/free_join/mod.rs

@@ -343,6 +343,13 @@ impl Database {
         &mut self.container_values
     }
 
+    pub fn restore_deserialized_runtime(&mut self) {
+        for (_, table) in self.tables.iter_mut() {
+            table.table.restore_deserialized_runtime();
+        }
+        self.container_values.restore_deserialized_runtime();
+    }
+
     pub fn rebuild_containers(&mut self, table_id: TableId) -> bool {
         let mut containers = mem::take(&mut self.container_values);
         let table = &self.tables[table_id].table;

core-relations/src/lib.rs

@@ -25,7 +25,9 @@ pub(crate) mod uf;
 mod tests;
 
 pub use action::{ExecutionState, MergeVal, QueryEntry, WriteVal};
-pub use base_values::{BaseValue, BaseValueId, BaseValuePrinter, BaseValues, Boxed};
+pub use base_values::{
+    register_deserializable_base_type, BaseValue, BaseValueId, BaseValuePrinter, BaseValues, Boxed,
+};
 pub use common::Value;
 pub use containers::{ContainerValue, ContainerValueId, ContainerValues};
 pub use free_join::{

core-relations/src/table/mod.rs

@@ -436,6 +436,37 @@ impl Table for SortedWritesTable {
     fn as_any_mut(&mut self) -> &mut dyn Any {
         self
     }
+
+    fn restore_deserialized_runtime(&mut self) {
+        self.subset_tracker.restore_deserialized_runtime();
+        self.hash = ShardedHashTable::default();
+        self.offsets.clear();
+
+        for (i, row) in self.data.data.iter().enumerate() {
+            if row[0].is_stale() {
+                continue;
+            }
+
+            let row_id = RowId::from_usize(i);
+            let (shard, hc) = hash_code(self.hash.shard_data(), row, self.n_keys);
+            self.hash.mut_shards()[shard.index()].insert_unique(
+                hc,
+                TableEntry {
+                    hashcode: hc as _,
+                    row: row_id,
+                },
+                TableEntry::hashcode,
+            );
+
+            if let Some(sort_by) = self.sort_by {
+                let sort_val = row[sort_by.index()];
+                if self.offsets.last().is_none_or(|(max, _)| sort_val > *max) {
+                    self.offsets.push((sort_val, row_id));
+                }
+            }
+        }
+    }
+
     fn clear(&mut self) {
         self.pending_state.clear();
         if self.data.data.len() == 0 {

core-relations/src/table_spec.rs

@@ -27,8 +27,7 @@ use crate::{
     offsets::{RowId, Subset, SubsetRef},
     pool::{with_pool_set, PoolSet, Pooled},
     row_buffer::{RowBuffer, TaggedRowBuffer},
-    DisplacedTable, DisplacedTableWithProvenance,
-    QueryEntry, TableId, Variable,
+    DisplacedTable, DisplacedTableWithProvenance, QueryEntry, TableId, Variable,
 };
 
 define_id!(pub ColumnId, u32, "a particular column in a table");
@@ -187,6 +186,13 @@ pub trait Table: Any + Send + Sync {
     /// A mutable variant of [`Table::as_any`] for downcasting.
     fn as_any_mut(&mut self) -> &mut dyn Any;
 
+    /// Restore runtime-only state after deserialization.
+    ///
+    /// Most tables serialize only durable data, so they do not need any extra
+    /// fixup and can use the default no-op implementation. Tables with
+    /// incremental caches or trackers should override this and reset them.
+    fn restore_deserialized_runtime(&mut self) {}
+
     /// The schema of the table.
     ///
     /// These are immutable properties of the table; callers can assume they
@@ -554,7 +560,9 @@ impl<'de> Deserialize<'de> for WrappedTable {
         } else if inner.as_any().is::<DisplacedTableWithProvenance>() {
             wrapper::<DisplacedTableWithProvenance>()
         } else {
-            return Err(serde::de::Error::custom("unknown table type for WrappedTable"));
+            return Err(serde::de::Error::custom(
+                "unknown table type for WrappedTable",
+            ));
         };
 
         Ok(WrappedTable { inner, wrapper })

egglog-bridge/src/lib.rs

@@ -1163,6 +1163,8 @@ impl EGraph {
     }
 
     pub fn restore_deserialized_runtime(&mut self) {
+        self.db.restore_deserialized_runtime();
+
         let funcs = self
             .funcs
             .iter()