wip, flattened type introduced

Signed-off-by: Kevin Mato <[email protected]>

Author: kvmto

runtime/common/ExecutionContext.h

@@ -17,6 +17,9 @@
 #include "cudaq/operators.h"
 #include <optional>
 #include <string_view>
+#include <iostream>
+
+#include "nvqir/stim/StimState.h"
 
 namespace cudaq {
 
@@ -27,18 +30,79 @@ struct RecordStorage{
   RecordStorage(size_t limit = 1e9) : memory_limit(limit), current_memory(0) {}
 
   std::vector<std::unique_ptr<SimulationState>> recordedStates;
-    void save_state(const SimulationState *state) {
-      recordedStates.push_back(std::make_unique<SimulationState>(*state););
+
+    void save_state(SimulationState *state) {
+        recordedStates.push_back(clone_state(state));
     }
-    std::vector<std::unique_ptr<SimulationState>> get_recorded_states() const { return recordedStates; }
+    const std::vector<std::unique_ptr<SimulationState>>& get_recorded_states() const { return recordedStates; }
 
     void clear() { recordedStates.clear(); }
     void dump_recorded_states() const {
       for (std::size_t i = 0; i < recordedStates.size(); i++) {
         recordedStates[i]->dump(std::cout);
       }
     }
-}
+    
+private:
+    std::unique_ptr<SimulationState> clone_state(SimulationState* state) {
+        if (state->isArrayLike()) {
+            // Handle array-like states (CusvState, etc.)
+            return clone_array_like_state(state);
+        } else {
+            // Handle specialized states (CuDensityMatState, StimState, etc.)
+            return clone_specialized_state(state);
+        }
+    }
+    
+    std::unique_ptr<SimulationState> clone_array_like_state(SimulationState* state) {
+        auto numQubits = state->getNumQubits();
+        if (numQubits > 20) {  // Prevent exponential explosion
+            throw std::runtime_error("State too large to clone via amplitudes");
+        }
+        
+        // Generate all basis states
+        auto totalStates = 1ULL << numQubits;
+        std::vector<std::vector<int>> basisStates;
+        for (size_t i = 0; i < totalStates; ++i) {
+            std::vector<int> basis(numQubits);
+            for (size_t j = 0; j < numQubits; ++j) {
+                basis[j] = (i >> j) & 1;
+            }
+            basisStates.push_back(basis);
+        }
+        
+        auto amplitudes = state->getAmplitudes(basisStates);
+        
+        // Create new state with appropriate precision
+        if (state->getPrecision() == SimulationState::precision::fp32) {
+            std::vector<std::complex<float>> floatAmps;
+            for (const auto& amp : amplitudes) {
+                floatAmps.emplace_back(static_cast<float>(amp.real()), 
+                                     static_cast<float>(amp.imag()));
+            }
+            return state->createFromData(floatAmps);
+        } else {
+            return state->createFromData(amplitudes);
+        }
+    }
+    
+    std::unique_ptr<SimulationState> clone_specialized_state(SimulationState* state) {
+        // Try dynamic_cast to known types that have clone methods
+        // this triggerd fatal error: library_types.h: No such file or directory
+        //if (auto* densityState = dynamic_cast<const CuDensityMatState*>(state)) {
+        //    return CuDensityMatState::clone(*densityState);
+        //}
+        
+        if (auto* stimState = dynamic_cast<const StimState*>(state)) {
+            return stimState->clone();
+        }
+        
+        // For unknown specialized types, try createFromSizeAndPtr as fallback
+        // This might work for some specialized states
+        // auto tensor = state->getTensor(0); 
+        // return state->createFromSizeAndPtr(tensor.get_num_elements(), tensor.data, 1);
+    }
+};
 
 /// The ExecutionContext is an abstraction to indicate how a CUDA-Q kernel
 /// should be executed.
@@ -179,12 +243,12 @@ class ExecutionContext {
   std::vector<std::vector<std::size_t>> errors_per_shot; // errors_per_shot[shot][error_id]
 
   /// @brief Save the current simulation state in the recorded states storage.
-  void save_state(const SimulationState *state){
+  void save_state(SimulationState *state){
     recordStorage.save_state(state);
   }
 
   /// @brief Get the recorded states saved during execution.
-  std::vector<std::unique_ptr<SimulationState>> get_recorded_states() const {
+  const std::vector<std::unique_ptr<SimulationState>>& get_recorded_states() const {
     return recordStorage.get_recorded_states();
   }
 

runtime/common/SimulationState.h

@@ -21,50 +21,16 @@ class SimulationState;
 /// Enum to specify the initial quantum state.
 enum class InitialState { ZERO, UNIFORM };
 
-struct StimData {
-    struct TableauClone {
-        std::vector<std::vector<bool>> x_output;
-        std::vector<std::vector<bool>> z_output;
-        std::vector<bool> r_output;
-
-        TableauClone copy() const {
-            TableauClone t;
-            t.x_output = x_output;
-            t.z_output = z_output;
-            t.r_output = r_output;
-            return t;
-        }
-    } tableau;
-
-    struct PauliFrameClone {
-        std::vector<bool> x;
-        std::vector<bool> z;
-
-        PauliFrameClone copy() const {
-            PauliFrameClone f;
-            f.x = x;
-            f.z = z;
-            return f;
-        }
-    } frame;
-
-    std::size_t current_size = 0;
-    std::size_t msm_err_count = 0;
-    uint64_t num_qubits = 0;
-    void* data() { return nullptr; }
-
-    // Copy helper for the entire StimData
-    StimData copy() const {
-        StimData s;
-        s.tableau = tableau.copy();
-        s.frame = frame.copy();
-        s.current_size = current_size;
-        s.msm_err_count = msm_err_count;
-        s.num_qubits = num_qubits;
-        return s;
-    }
-};
 
+/// @brief StimData now stores a list of (pointer, size) pairs
+/// according to convention:
+/// 0: pointer to num_qubits, size = 1
+/// 1: pointer to msm_err_count, size = 1
+/// 2: pointer to num_stabilizers, size = 1
+/// 3: x_output array, size = x_output_size (X stabilizers + destabilisers + phase bits)
+/// 4: z_output array, size = z_output_size (Z stabilizers + destabilisers + phase bits)
+/// 5: frame array (Pauli frame), size = 2*num_qubits
+using StimData = std::vector<std::pair<void*, std::size_t>>;
 
 /// @brief Encapsulates a list of tensors (data pointer and dimensions).
 // Note: tensor data is expected in column-major.
@@ -183,7 +149,7 @@ class SimulationState {
             "Cannot initialize state vector/density matrix state by stim "
             "data. Please use stabilizer simulator backends.");
       auto &dataCasted = std::get<StimData>(data);
-      return createFromSizeAndPtr(1, const_cast<StimData*>(&dataCasted), data.index());
+      return createFromSizeAndPtr(dataCasted.size(), const_cast<StimData*>(&dataCasted), data.index());
     }
     // Flat array state data
     // Check the precision first. Get the size and