StructRef: is it possible to have a literal type?

Numba Community Support
1

In some code I’m writing I would be interested in using StructRefs with some literal type.
Essentially, to de-facto overload/dispatch to differently-compiled methods of the same function depending on the value of a field in the structure.

This would mimic Julia’s Type Parameters, which can be any valid is bits type, and is a rather handy feature.

Is this possible?

2

Hi @PhilipVinc

Are you looking to do this sort of thing?

import numpy as np

from numba import njit
from numba.core import types
from numba.experimental import structref
from numba.extending import overload

@structref.register
class MyStructType(types.StructRef):
    def preprocess_fields(self, fields):
        for name, typ in fields:
            assert isinstance(typ, types.IntegerLiteral)
        return tuple((name, typ) for name, typ in fields)

class MyStruct(structref.StructRefProxy):
    def __new__(cls, value):
        return structref.StructRefProxy.__new__(cls, value)

    @property
    def value(self):
        return MyStruct_get_value(self)


@njit
def MyStruct_get_value(self):
    return self.value

structref.define_proxy(MyStruct, MyStructType, ["value"])

def bar(x):
    pass

@overload(bar)
def ol_bar(x):
    name, typ = x._fields[0]
    if typ.literal_value < 7:
        def impl(x):
            print(".lt. 7")
    else:
        def impl(x):
            print(".ge. 7")
    return impl


@njit
def foo():
    struct10 = MyStruct(10)
    bar(struct10)
    struct2 = MyStruct(2)
    bar(struct2)

foo()

Hope this helps?