Skip to content

Commit 0714c31

Browse files
markhalkamarkhalka
committed
fix uninitialized enum
Signed-off-by: Mark Halka <Mark.Halka@Point72.com>
1 parent c78da59 commit 0714c31

3 files changed

Lines changed: 116 additions & 66 deletions

File tree

‎cpp/csp/python/PyNumbaNode.cpp‎

Lines changed: 71 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#include <csp/engine/InputId.h>
22
#include <csp/engine/Node.h>
33
#include <csp/engine/CspEnum.h>
4+
#include <csp/engine/PartialSwitchCspType.h>
45
#include <csp/engine/Struct.h>
56
#include <csp/python/Conversions.h>
67
#include <csp/python/Exception.h>
@@ -17,6 +18,7 @@
1718
#include <unordered_map>
1819
#include <cstring>
1920
#include <iostream>
21+
#include <type_traits>
2022

2123
namespace csp::python
2224
{
@@ -30,18 +32,54 @@ extern "C" CSPIMPL_EXPORT int64_t csp_numba_struct_enum_field_value( const void
3032
return enumField -> value();
3133
}
3234

33-
extern "C" CSPIMPL_EXPORT void csp_numba_struct_enum_field_set( void * struct_ptr, int64_t field_offset, int64_t value )
35+
extern "C" CSPIMPL_EXPORT void csp_numba_struct_enum_field_set(
36+
void * struct_ptr,
37+
int64_t field_offset,
38+
int64_t value,
39+
const void * struct_type
40+
)
3441
{
35-
// Rebuild the field from its existing enum meta rather than treating the
36-
// destination as a raw int64 slot.
37-
char * bytes = static_cast<char *>( struct_ptr );
38-
auto * enumField = reinterpret_cast<CspEnum *>( bytes + field_offset );
39-
*enumField = enumField -> meta() -> create( value );
42+
// An unset enum field does not contain enum metadata. Resolve the declared
43+
// field type through the owning struct type and use the field setter so its
44+
// set bit is updated as well.
45+
const auto * pyStructMeta = static_cast<const PyStructMeta *>( struct_type );
46+
const auto & fields = pyStructMeta -> structMeta -> fields();
47+
auto fieldIt = std::find_if( fields.begin(), fields.end(), [field_offset]( const auto & field )
48+
{
49+
return field -> offset() == static_cast<size_t>( field_offset );
50+
} );
51+
CSP_ASSERT( fieldIt != fields.end() );
52+
CSP_ASSERT( ( *fieldIt ) -> type() -> type() == CspType::Type::ENUM );
53+
54+
const auto * enumType = static_cast<const CspEnumType *>( ( *fieldIt ) -> type().get() );
55+
( *fieldIt ) -> setValue<CspEnum>(
56+
static_cast<Struct *>( struct_ptr ),
57+
enumType -> meta() -> create( value )
58+
);
4059
}
4160

4261
namespace
4362
{
4463
constexpr size_t OUTPUT_VALUE_SLOT_BYTES = sizeof( int64_t );
64+
65+
using NumbaInputTypeSwitch = PartialSwitchCspType<
66+
CspType::Type::BOOL,
67+
CspType::Type::INT64,
68+
CspType::Type::DOUBLE,
69+
CspType::Type::DATETIME,
70+
CspType::Type::TIMEDELTA,
71+
CspType::Type::ENUM,
72+
CspType::Type::STRUCT
73+
>;
74+
75+
using NumbaOutputTypeSwitch = PartialSwitchCspType<
76+
CspType::Type::BOOL,
77+
CspType::Type::INT64,
78+
CspType::Type::DOUBLE,
79+
CspType::Type::DATETIME,
80+
CspType::Type::TIMEDELTA,
81+
CspType::Type::ENUM
82+
>;
4583
}
4684

4785
PyNumbaNode::PyNumbaNode(
@@ -303,33 +341,21 @@ void PyNumbaNode::executeImpl()
303341
{
304342
const CspType * cspType = ts -> type();
305343
CSP_ASSERT( cspType != nullptr );
306-
switch( cspType -> type() )
344+
NumbaInputTypeSwitch::invoke( cspType, [this, ts, i]( auto tag )
307345
{
308-
case CspType::Type::INT64:
309-
m_inputArgs[ i ] = &( ts -> lastValueTyped<int64_t>() );
310-
break;
311-
case CspType::Type::DOUBLE:
312-
m_inputArgs[ i ] = &( ts -> lastValueTyped<double>() );
313-
break;
314-
case CspType::Type::BOOL:
315-
m_inputArgs[ i ] = &( ts -> lastValueTyped<bool>() );
316-
break;
317-
case CspType::Type::DATETIME:
318-
m_inputArgs[ i ] = &( ts -> lastValueTyped<DateTime>() );
319-
break;
320-
case CspType::Type::TIMEDELTA:
321-
m_inputArgs[ i ] = &( ts -> lastValueTyped<TimeDelta>() );
322-
break;
323-
case CspType::Type::ENUM:
346+
using CType = typename decltype( tag )::type;
347+
if constexpr( std::is_same_v<CType, CspEnum> )
348+
{
324349
m_inputEnumStorage[ i ] = ts -> lastValueTyped<CspEnum>().value();
325350
m_inputArgs[ i ] = &m_inputEnumStorage[ i ];
326-
break;
327-
case CspType::Type::STRUCT:
351+
}
352+
else if constexpr( std::is_same_v<CType, StructPtr> )
353+
{
328354
m_inputArgs[ i ] = ts -> lastValueTyped<StructPtr>().get();
329-
break;
330-
default:
331-
break;
332-
}
355+
}
356+
else
357+
m_inputArgs[ i ] = &( ts -> lastValueTyped<CType>() );
358+
} );
333359
}
334360
}
335361
}
@@ -360,52 +386,38 @@ void PyNumbaNode::executeImpl()
360386
{
361387
const CspType * cspType = ts -> type();
362388
CSP_ASSERT( cspType != nullptr );
363-
switch( cspType -> type() )
389+
NumbaOutputTypeSwitch::invoke( cspType, [this, ts, cspType, i, currentCycleCount, currentTime]( auto tag )
364390
{
365-
case CspType::Type::INT64:
366-
ts -> outputTickTyped<int64_t>(
367-
currentCycleCount, currentTime,
368-
*static_cast<int64_t *>( m_outputValueSlots[ i ] )
369-
);
370-
break;
371-
case CspType::Type::DOUBLE:
372-
ts -> outputTickTyped<double>(
373-
currentCycleCount, currentTime,
374-
*static_cast<double *>( m_outputValueSlots[ i ] )
375-
);
376-
break;
377-
case CspType::Type::BOOL:
378-
ts -> outputTickTyped<bool>(
379-
currentCycleCount, currentTime,
380-
*static_cast<bool *>( m_outputValueSlots[ i ] )
381-
);
382-
break;
383-
case CspType::Type::DATETIME:
391+
using CType = typename decltype( tag )::type;
392+
if constexpr( std::is_same_v<CType, DateTime> )
393+
{
384394
ts -> outputTickTyped<DateTime>(
385395
currentCycleCount, currentTime,
386396
DateTime::fromNanoseconds( *static_cast<int64_t *>( m_outputValueSlots[ i ] ) )
387397
);
388-
break;
389-
case CspType::Type::TIMEDELTA:
398+
}
399+
else if constexpr( std::is_same_v<CType, TimeDelta> )
400+
{
390401
ts -> outputTickTyped<TimeDelta>(
391402
currentCycleCount, currentTime,
392403
TimeDelta::fromNanoseconds( *static_cast<int64_t *>( m_outputValueSlots[ i ] ) )
393404
);
394-
break;
395-
case CspType::Type::ENUM:
405+
}
406+
else if constexpr( std::is_same_v<CType, CspEnum> )
396407
{
397408
auto enumType = static_cast<const CspEnumType *>( cspType );
398409
int64_t enumValue = *static_cast<int64_t *>( m_outputValueSlots[ i ] );
399-
CspEnum enumInstance = enumType -> meta() -> create( enumValue );
400410
ts -> outputTickTyped<CspEnum>(
401411
currentCycleCount, currentTime,
402-
enumInstance
412+
enumType -> meta() -> create( enumValue )
403413
);
404-
break;
405414
}
406-
default:
407-
break;
408-
}
415+
else
416+
ts -> outputTickTyped<CType>(
417+
currentCycleCount, currentTime,
418+
*static_cast<CType *>( m_outputValueSlots[ i ] )
419+
);
420+
} );
409421
}
410422
}
411423
}

‎csp/impl/wiring/csp_numba/struct_support.py‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -63,28 +63,35 @@ def codegen(context, builder, signature, args):
6363

6464

6565
@intrinsic
66-
def struct_enum_store(typingctx, struct_ptr, field_offset_const, value):
66+
def struct_enum_store(typingctx, struct_ptr, field_offset_const, value, struct_type_address_const):
6767
"""Store an integer enum value into a CspEnum struct field."""
68-
if struct_ptr == numba_types.voidptr and isinstance(field_offset_const, numba_types.Literal):
68+
if (
69+
struct_ptr == numba_types.voidptr
70+
and isinstance(field_offset_const, numba_types.Literal)
71+
and isinstance(struct_type_address_const, numba_types.Literal)
72+
):
6973
field_offset = field_offset_const.literal_value
70-
sig = numba_types.void(struct_ptr, field_offset_const, numba_types.int64)
74+
struct_type_address = struct_type_address_const.literal_value
75+
sig = numba_types.void(struct_ptr, field_offset_const, numba_types.int64, struct_type_address_const)
7176

7277
def codegen(context, builder, signature, args):
73-
[struct_ptr_val, _field_offset, enum_value] = args
78+
[struct_ptr_val, _field_offset, enum_value, _struct_type_address] = args
7479

7580
i8p = ir.IntType(8).as_pointer()
7681
i64 = ir.IntType(64)
7782
module = builder.module
7883
fn_name = "csp_numba_struct_enum_field_set"
7984
fn = module.globals.get(fn_name)
8085
if fn is None:
81-
fn_ty = ir.FunctionType(ir.VoidType(), [i8p, i64, i64])
86+
fn_ty = ir.FunctionType(ir.VoidType(), [i8p, i64, i64, i8p])
8287
fn = ir.Function(module, fn_ty, name=fn_name)
8388
fn.attributes.add("nounwind")
8489

8590
struct_bytes = builder.bitcast(struct_ptr_val, i8p)
8691
offset_val = context.get_constant(numba_types.int64, field_offset)
87-
builder.call(fn, [struct_bytes, offset_val, enum_value])
92+
struct_type_address_val = context.get_constant(numba_types.uintp, struct_type_address)
93+
struct_type_ptr = builder.inttoptr(struct_type_address_val, i8p)
94+
builder.call(fn, [struct_bytes, offset_val, enum_value, struct_type_ptr])
8895
return context.get_dummy_value()
8996

9097
return sig, codegen
@@ -159,7 +166,13 @@ def set_field(self, struct_name: str, field_name: str, value_expr: ast.AST) -> a
159166
# reconstructed back into the field's native CspEnum representation.
160167
field_info = self._get_field_info(field_name)
161168
struct_ptr = ast.Name(id=struct_name, ctx=ast.Load())
162-
return AST.function_call("struct_enum_store", struct_ptr, ast.Constant(value=field_info.offset), value_expr)
169+
return AST.function_call(
170+
"struct_enum_store",
171+
struct_ptr,
172+
ast.Constant(value=field_info.offset),
173+
value_expr,
174+
ast.Constant(value=id(self.value)),
175+
)
163176

164177
@classmethod
165178
def try_parse_state(cls, node: ast.AnnAssign, var_name: str, globalns: dict) -> Optional[StateVariableInfo]:

‎csp/tests/test_numba_node.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1038,6 +1038,31 @@ def g():
10381038

10391039

10401040
class TestStructSupport(unittest.TestCase):
1041+
def test_assign_unset_struct_enum_field(self):
1042+
class Direction(csp.Enum):
1043+
UP = 1
1044+
DOWN = -1
1045+
1046+
class Velocity(csp.Struct):
1047+
direction: Direction
1048+
1049+
@numba_node
1050+
def set_direction(x: ts[int]) -> ts[int]:
1051+
with csp.state():
1052+
s_velocity: Velocity = Velocity()
1053+
1054+
s_velocity.direction = Direction.UP
1055+
if s_velocity.direction == Direction.UP:
1056+
return x
1057+
return 0
1058+
1059+
@csp.graph
1060+
def g():
1061+
csp.add_graph_output("result", set_direction(csp.const(7)))
1062+
1063+
results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1))
1064+
self.assertEqual([value for _, value in results["result"]], [7])
1065+
10411066
def test_invalid_struct_output(self):
10421067
class Point(csp.Struct):
10431068
x: float

0 commit comments

Comments
 (0)