Handle prd updates for a state we already know about but don't have the keys

src/api.rs

@@ -847,11 +847,28 @@ fn handle_prd(
         PrdType::Update => {
             // Compute the merkle tree root for the proposed new state to see if we already know about it
             let update_merkle_root = prd.pcd_commitments.create_merkle_tree()?.root().ok_or(AnyhowError::msg("Invalid merkle tree"))?;
-            if relevant_process.get_state_for_id(&update_merkle_root).is_ok() {
+            let updated_state: &ProcessState = if let Ok(existing_state) = relevant_process.get_state_for_id_mut(&update_merkle_root) {
-                // We already know about that state
+                // We already know about that state, if we also have the keys we can just stop here
-                return Err(AnyhowError::msg("Received update for a state we already know"));
+                if !existing_state.keys.is_empty() {
+                    // We check that the keys are the same, just in case
+                    if existing_state.keys == prd.keys {
+                        debug!("Received update for a state we already know and have the same keys");
+                        return Ok(ApiReturn::default());
+                    } else {
+                        // We merge the keys that we have with the ones that we receive
+                        debug!("Received update for a state we already know and have different keys, merging them");
+                        for (key, value) in prd.keys.iter() {
+                            if !existing_state.keys.contains_key(key) { // We don't want to override existing keys
+                                existing_state.keys.insert(key.clone(), value.clone());
                             }
-
+                        }
+                    }
+                } else {
+                    // We don't have any keys for this state, we can just update it
+                    existing_state.keys = prd.keys.clone();
+                }
+                existing_state
+            } else {
                 let commited_in = relevant_process.get_process_tip()?;
 
                 let new_state = ProcessState {
@@ -864,11 +881,14 @@ fn handle_prd(
                     ..Default::default()
                 };
 
-            // Compute the diffs
-            let diffs = create_diffs(&lock_local_device()?, &relevant_process, &new_state, members_list)?;
-
                 relevant_process.insert_concurrent_state(new_state)?;
 
+                relevant_process.get_state_for_id(&new_state.state_id).expect("New state should be inserted")
+            };
+
+            // Compute the diffs
+            let diffs = create_diffs(&lock_local_device()?, &relevant_process, &updated_state, members_list)?;
+
             let updated_process = UpdatedProcess {
                 process_id: outpoint,
                 current_process: relevant_process.clone(),