[Feature Request] Adding a drjitstruct decorator

[DoeringChristian]

Hi,
Working with the DRJIT_STRUCT feature in python, I have found that It would be useful to have a drjitstruct decorator.
This decorator would generate the required DRJIT_STRUCT dictionary from the annotations of the class variables.

If I want to declare a class as a struct for Dr.Jit I would have to add the DRJIT_STRUCT dict to the class and re-define the variables. This is somewhat redundant and also could lead to potential problems when one forgets to specify some variables.

class RestirSample:
    x_v: mi.Vector3f
    n_v: mi.Vector3f
    x_s: mi.Vector3f
    n_s: mi.Vector3f

    L_o: mi.Color3f
    p_q: mi.Float
    f: mi.Color3f
    valid: mi.Bool

    DRJIT_STRUCT = {
        "x_v": mi.Vector3f,
        "n_v": mi.Vector3f,
        "x_s": mi.Vector3f,
        "n_s": mi.Vector3f,
        "L_o": mi.Color3f,
        "p_q": mi.Float,
        "f": mi.Color3f,
        "valid": mi.Bool,
    }

My proposal would be to add a drjitstruct decorator that generates this dictionary automatically.

@drjitstruct
class RestirSample:
    x_v: mi.Vector3f
    n_v: mi.Vector3f
    x_s: mi.Vector3f
    n_s: mi.Vector3f

    L_o: mi.Color3f
    p_q: mi.Float
    f: mi.Color3f
    valid: mi.Bool

A potential way of implementing such a decorator would be the following.

from typing import get_type_hints

def drjitstruct(cls):
    drjit_struct = {}

    type_hints = get_type_hints(cls)

    for name, ty in type_hints.items():
        drjit_struct[name] = ty
    cls.DRJIT_STRUCT = drjit_struct
    return cls

I have used this implementation in some of my projects and wanted to ask iff there is interest in implementing such a feature directly in Dr.Jit before opening a PR.

[njroussel]

Hi @DoeringChristian

This is something we'd like to have. I can't exactly remember the details, but we had attempted something similar to your suggested implementation and found some corner cases which would not be handled.

Thanks for expressing interest in this, we'll have a look at it again 👍