feat(storage/state_update): store re-declared Cairo 0 class information

This change adds a new `redeclared_classes` table storing block numbers
at which class re-declarations did happen.

When inserting a state update we take care of adding new rows here if
the state update contains a re-declaration -- and then when retrieving
the state update from storage we add re-declared classes to the set
of Cairo classes declared at the block.

crates/storage/src/connection/state_update.rs

@@ -120,10 +120,18 @@ impl Transaction<'_> {
             .prepare_cached(
                 r"INSERT INTO class_definitions (block_number, hash) VALUES (?1, ?2)
                 ON CONFLICT(hash)
-                    DO UPDATE SET block_number=excluded.block_number
-                    WHERE block_number IS NULL",
+                    DO UPDATE SET block_number=IFNULL(block_number,excluded.block_number)
+                RETURNING block_number",
             )
             .context("Preparing class hash and block number upsert statement")?;
+
+        let mut insert_redeclared_class = self
+            .inner()
+            .prepare_cached(
+                r"INSERT INTO redeclared_classes (class_hash, block_number) VALUES (?, ?)",
+            )
+            .context("Preparing redeclared class insert statement")?;
+
         // OR IGNORE is required to handle legacy syncing logic, where the casm
         // definition is inserted before the state update
         let mut insert_casm_hash = self
@@ -212,6 +220,20 @@ impl Transaction<'_> {
             .keys()
             .map(|sierra| ClassHash(sierra.0));
         let cairo = declared_cairo_classes.iter().copied();
+
+        let declared_classes = sierra.chain(cairo);
+
+        for class in declared_classes {
+            let declared_at = upsert_declared_at
+                .query_row(params![&block_number, &class], |row| {
+                    row.get_block_number(0)
+                })?;
+            if declared_at != block_number {
+                tracing::debug!(%declared_at, %block_number, class_hash=%class, "Re-declared class");
+                insert_redeclared_class.execute(params![&class, &block_number])?;
+            }
+        }
+
         // Older cairo 0 classes were never declared, but instead got implicitly
         // declared on first deployment. Until such classes disappear we need to
         // cater for them here. This works because the sql only updates the row
@@ -223,10 +245,10 @@ impl Transaction<'_> {
                 _ => None,
             });
 
-        let declared_classes = sierra.chain(cairo).chain(deployed);
-
-        for class in declared_classes {
-            upsert_declared_at.execute(params![&block_number, &class])?;
+        for class in deployed {
+            let _ = upsert_declared_at.query_row(params![&block_number, &class], |row| {
+                row.get_block_number(0)
+            })?;
         }
 
         for (sierra_hash, casm_hash) in declared_sierra_classes {
@@ -393,6 +415,22 @@ impl Transaction<'_> {
             };
         }
 
+        let mut stmt = self
+            .inner()
+            .prepare_cached(r"SELECT class_hash FROM redeclared_classes WHERE block_number = ?")
+            .context("Preparing re-declared class query statement")?;
+
+        let mut redeclared_classes = stmt
+            .query_map(params![&block_number], |row| row.get_class_hash(0))
+            .context("Querying re-declared classes")?;
+        while let Some(class_hash) = redeclared_classes
+            .next()
+            .transpose()
+            .context("Iterating over re-declared classes")?
+        {
+            state_update = state_update.with_declared_cairo_class(class_hash);
+        }
+
         let mut stmt = self
         .inner().prepare_cached(
             r"SELECT
@@ -480,18 +518,17 @@ impl Transaction<'_> {
         let mut stmt = self
             .inner()
             .prepare_cached(
-                r"
-                SELECT COUNT(block_number)
-                FROM canonical_blocks
-                LEFT JOIN class_definitions ON canonical_blocks.number = class_definitions.block_number
-                WHERE number >= ?
-                GROUP BY canonical_blocks.number
-                ORDER BY canonical_blocks.number ASC
+                r"SELECT
+                  (SELECT COUNT(block_number) FROM class_definitions WHERE block_number=block_headers.number) +
+                  (SELECT COUNT(block_number) FROM redeclared_classes WHERE block_number=block_headers.number)
+                FROM block_headers
+                WHERE block_headers.number >= ?
+                ORDER BY block_headers.number ASC
                 LIMIT ?",
             )
             .context("Preparing get number of declared classes statement")?;
 
-        let max_len = u64::try_from(max_num_blocks.get()).expect("ptr size is 64 bits");
+        let max_len: u64 = u64::try_from(max_num_blocks.get()).expect("ptr size is 64 bits");
         let mut counts = stmt
             .query_map(params![&start, &max_len], |row| row.get(0))
             .context("Querying declared classes counts")?;
@@ -533,6 +570,23 @@ impl Transaction<'_> {
             result.push(class_hash);
         }
 
+        let mut stmt = self
+            .inner()
+            .prepare_cached(r"SELECT class_hash FROM redeclared_classes WHERE block_number = ?")
+            .context("Preparing re-declared class query")?;
+
+        let mut redeclared_classes = stmt
+            .query_map(params![&block_number], |row| row.get_class_hash(0))
+            .context("Querying re-declared classes")?;
+
+        while let Some(class_hash) = redeclared_classes
+            .next()
+            .transpose()
+            .context("Iterating over re-declared classes")?
+        {
+            result.push(class_hash)
+        }
+
         Ok(Some(result))
     }
 
@@ -1214,6 +1268,34 @@ mod tests {
             assert_eq!(non_existent, None);
         }
 
+        #[test]
+        fn redeclared_classes() {
+            let (mut db, _state_update, header) = setup();
+
+            let tx = db.transaction().unwrap();
+            let new_header = header
+                .child_builder()
+                .finalize_with_hash(block_hash!("0xabcdee"));
+            let new_state_update = StateUpdate::default()
+                .with_block_hash(new_header.hash)
+                .with_declared_cairo_class(CAIRO_HASH);
+            tx.insert_block_header(&new_header).unwrap();
+            tx.insert_state_update(new_header.number, &new_state_update)
+                .unwrap();
+
+            tx.commit().unwrap();
+
+            let tx = db.transaction().unwrap();
+
+            let result = tx.state_update(new_header.number.into()).unwrap().unwrap();
+            assert_eq!(result, new_state_update);
+
+            let result = tx
+                .declared_classes_counts(new_header.number.into(), NonZeroUsize::new(1).unwrap())
+                .unwrap();
+            assert_eq!(result[0], 1);
+        }
+
         #[test]
         fn reverse_state_update() {
             let (mut db, _state_update, header) = setup();

crates/storage/src/schema.rs

@@ -23,6 +23,7 @@ mod revision_0060;
 mod revision_0061;
 mod revision_0062;
 mod revision_0063;
+mod revision_0064;
 
 pub(crate) use base::base_schema;
 
@@ -54,6 +55,7 @@ pub fn migrations() -> &'static [MigrationFn] {
         revision_0061::migrate,
         revision_0062::migrate,
         revision_0063::migrate,
+        revision_0064::migrate,
     ]
 }
 

crates/storage/src/schema/revision_0064.rs

@@ -0,0 +1,14 @@
+use anyhow::Context;
+
+pub(crate) fn migrate(tx: &rusqlite::Transaction<'_>) -> anyhow::Result<()> {
+    tx.execute_batch(
+        r"CREATE TABLE redeclared_classes (
+            class_hash BLOB NOT NULL,
+            block_number INTEGER NOT NULL REFERENCES block_headers(number) ON DELETE CASCADE
+        );
+        CREATE INDEX redeclared_classes_block_number ON redeclared_classes(block_number);",
+    )
+    .context("Adding redeclared_classes table")?;
+
+    Ok(())
+}