fix container memory cleanup

Signed-off-by: Mark Halka <mark.halka2001@gmail.com>

Author: markhalka

cpp/csp/python/PyNumbaNode.cpp

@@ -75,7 +75,8 @@ PyNumbaNode::~PyNumbaNode()
 
     for( size_t i = 0; i < m_stateCount; ++i )
     {
-        if( m_stateArgs[ i ] != nullptr )
+        if( m_stateArgs[ i ] != nullptr &&
+            m_containerStateIndices.find( i ) == m_containerStateIndices.end() )
             delete[] static_cast<char *>( m_stateArgs[ i ] );
     }
     delete[] m_stateArgs;
@@ -168,6 +169,7 @@ void PyNumbaNode::initStateArrays( PyObjectPtr stateVariables, PyObjectPtr nrtSt
             {
                 size_t idxVal = static_cast<size_t>( PyLong_AsLongLong( idx ) );
                 nrtSet.insert( idxVal );
+                m_containerStateIndices.insert( idxVal );
             }
         }
     }

cpp/csp/python/PyNumbaNode.h

@@ -6,6 +6,7 @@
 #include <csp/python/PyObjectPtr.h>
 #include <Python.h>
 #include <cstdint>
+#include <unordered_set>
 
 namespace csp::python
 {
@@ -103,6 +104,7 @@ class PyNumbaNode final : public csp::Node
     // State arrays
     void ** m_stateArgs = nullptr;
     size_t m_stateCount = 0;
+    std::unordered_set<size_t> m_containerStateIndices;
 
     // Keep Python objects alive
     PyObjectPtr m_dataReference;

csp/impl/wiring/csp_numba/input_handlers.py

@@ -8,6 +8,7 @@
 """
 
 import inspect
+from dataclasses import dataclass
 from typing import Any, Optional, get_args, get_origin
 
 from numba_cfunc_compiler.function_analyzer import InputTypeHandler
@@ -41,28 +42,43 @@ def _parse_list_basket_annotation(ann: Any) -> tuple[bool, Optional[type]]:
     return False, None
 
 
-def _parse_dict_basket_annotation(ann: Any) -> tuple[bool, Optional[type]]:
+@dataclass(frozen=True)
+class _DictBasketExpectedType:
+    key_type: type
+    element_type: type
+
+
+def _parse_dict_basket_annotation(ann: Any) -> tuple[bool, Optional[type], Optional[type]]:
     key_type = None
     value_ann = None
 
     # Legacy syntax: {str: ts[T]} / {int: ts[T]}
     if isinstance(ann, dict):
         if len(ann) != 1:
-            return True, None
+            return True, None, None
         key_type, value_ann = next(iter(ann.items()))
     # Modern syntax: dict[str, ts[T]] / typing.Dict[str, ts[T]]
     elif get_origin(ann) is dict:
         args = get_args(ann)
         if len(args) != 2:
-            return True, None
+            return True, None, None
         key_type, value_ann = args
     else:
-        return False, None
+        return False, None, None
 
     if key_type not in (str, int):
-        return False, None
+        return False, None, None
 
-    return True, _extract_ts_inner_type(value_ann)
+    return True, key_type, _extract_ts_inner_type(value_ann)
+
+
+def _validate_edge_type(param_name: str, edge: Edge, expected_type: type) -> Edge:
+    actual_type = edge.tstype.typ
+    if actual_type is not expected_type and not (expected_type is float and actual_type is int):
+        raise TypeError(
+            f"Argument '{param_name}' expected ts[{expected_type.__name__}], got ts[{actual_type.__name__}]"
+        )
+    return edge
 
 
 class TsInputHandler(InputTypeHandler):
@@ -78,7 +94,7 @@ def try_parse(self, param: inspect.Parameter, ann: Any) -> Optional[ParameterInf
     def validate_value(self, param_name: str, value: Any, expected_type: Any) -> Any:
         if not isinstance(value, Edge):
             raise TypeError(f"Argument '{param_name}' must be an Edge, got {type(value).__name__}")
-        return value
+        return _validate_edge_type(param_name, value, expected_type)
 
 
 class ListBasketInputHandler(InputTypeHandler):
@@ -96,36 +112,48 @@ def try_parse(self, param: inspect.Parameter, ann: Any) -> Optional[ParameterInf
     def validate_value(self, param_name: str, value: Any, expected_type: Any) -> Any:
         if not isinstance(value, (list, tuple)):
             raise TypeError(f"Argument '{param_name}' must be a list, got {type(value).__name__}")
+        if not value:
+            raise ValueError(f"List basket '{param_name}' must not be empty")
 
         result = {}
         for i, edge in enumerate(value):
             if not isinstance(edge, Edge):
                 raise TypeError(f"List basket '{param_name}[{i}]' must be an Edge, got {type(edge).__name__}")
-            result[i] = edge
+            result[i] = _validate_edge_type(f"{param_name}[{i}]", edge, expected_type)
         return result
 
 
 class DictBasketInputHandler(InputTypeHandler):
     """Handles {key_type: ts[type]} and dict[key_type, ts[type]] basket annotations."""
 
     def try_parse(self, param: inspect.Parameter, ann: Any) -> Optional[ParameterInfo]:
-        matched, inner_type = _parse_dict_basket_annotation(ann)
+        matched, key_type, inner_type = _parse_dict_basket_annotation(ann)
         if not matched:
             return None
         if inner_type is None:
             raise TypeError(f"Dict basket '{param.name}' element ts[type] is missing type argument")
 
-        return ParameterInfo(expected_type=inner_type, category="signal_set")
+        return ParameterInfo(
+            expected_type=_DictBasketExpectedType(key_type=key_type, element_type=inner_type),
+            category="signal_set",
+        )
 
     def validate_value(self, param_name: str, value: Any, expected_type: Any) -> Any:
         if not isinstance(value, dict):
             raise TypeError(f"Argument '{param_name}' must be a dict, got {type(value).__name__}")
+        if not value:
+            raise ValueError(f"Dict basket '{param_name}' must not be empty")
 
         result = {}
         for key, edge in value.items():
+            if not isinstance(key, expected_type.key_type):
+                raise TypeError(
+                    f"Dict basket '{param_name}' key {key!r} must be "
+                    f"{expected_type.key_type.__name__}, got {type(key).__name__}"
+                )
             if not isinstance(edge, Edge):
                 raise TypeError(f"Dict basket '{param_name}[{key!r}]' must be an Edge, got {type(edge).__name__}")
-            result[key] = edge
+            result[key] = _validate_edge_type(f"{param_name}[{key!r}]", edge, expected_type.element_type)
         return result
 
 

csp/impl/wiring/csp_numba/signal_set_support.py

@@ -332,7 +332,7 @@ def handle_signal_set_iteration(converter, node: ast.For) -> Optional[ast.AST]:
     else:
         filter_array = None
 
-    iter_var = f"_iter_{signal_set_name}"
+    iter_var = f"_iter_{signal_set_name}_{converter.variable_factory.create_temporary_variable_name()}"
     range_call = ast.Call(
         func=ast.Name(id="range", ctx=ast.Load()),
         args=[ast.Constant(value=signal_set.length)],

csp/tests/test_numba_node.py

@@ -541,6 +541,54 @@ def g():
 
 
 class TestFeatures(unittest.TestCase):
+    def test_rejects_mismatched_signal_type(self):
+        @numba_node
+        def typed_signal(x: ts[int]) -> ts[int]:
+            return 1
+
+        @csp.graph
+        def g():
+            csp.add_graph_output("result", typed_signal(csp.const(1.25)))
+
+        with self.assertRaisesRegex(TypeError, r"expected ts\[int\], got ts\[float\]"):
+            csp.build_graph(g)
+
+    def test_allows_int_to_float_signal_upcast(self):
+        @numba_node
+        def typed_signal(x: ts[float]) -> ts[float]:
+            return x + 0.5
+
+        @csp.graph
+        def g():
+            csp.add_graph_output("result", typed_signal(csp.const(1)))
+
+        results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1))
+        self.assertEqual([v for _, v in results["result"]], [1.5])
+
+    def test_rejects_mismatched_basket_element_type(self):
+        @numba_node
+        def typed_basket(xs: list[ts[int]]) -> ts[int]:
+            return 1
+
+        @csp.graph
+        def g():
+            csp.add_graph_output("result", typed_basket([csp.const(1), csp.const(2.5)]))
+
+        with self.assertRaisesRegex(TypeError, r"expected ts\[int\], got ts\[float\]"):
+            csp.build_graph(g)
+
+    def test_rejects_mismatched_dict_basket_key_type(self):
+        @numba_node
+        def typed_basket(xs: dict[int, ts[int]]) -> ts[int]:
+            return 1
+
+        @csp.graph
+        def g():
+            csp.add_graph_output("result", typed_basket({"bad": csp.const(1)}))
+
+        with self.assertRaisesRegex(TypeError, r"key 'bad' must be int, got str"):
+            csp.build_graph(g)
+
     def test_lifecycle_start(self):
         @numba_node
         def with_start(x: ts[int]) -> ts[int]:
@@ -796,6 +844,23 @@ def g():
         self.assertEqual([v for _, v in results["tick_count"]], [1, 2])
         self.assertEqual([v for _, v in results["all_valid"]], [0, 1])
 
+    def test_nested_iteration_over_same_basket(self):
+        @numba_node
+        def cross_product(xs: list[ts[int]]) -> ts[int]:
+            total = 0
+            for i in xs.keys():
+                for j in xs.keys():
+                    total = total + xs[i] * xs[j]
+            return total
+
+        @csp.graph
+        def g():
+            result = cross_product([csp.const(2), csp.const(3)])
+            csp.add_graph_output("result", result)
+
+        results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1))
+        self.assertEqual([v for _, v in results["result"]], [25])
+
     def test_multiple_outputs(self):
         @numba_node
         def multi_out(x: ts[int]) -> csp.Outputs(doubled=ts[int], squared=ts[int], positive=ts[int]):