Need syntax for passing `CPtr`s to structs and arrays of structs

[dylon]

The following works perfectly well in LPython and LPython generates valid C code from it, but the syntax is not supported by CPython:

# file: main.py
from lpython import i32, dataclass

from numpy import empty

@dataclass
class Foo:
     x: i32
     y: i32

def init(foos: Foo[:]) -> None:
    foos[0] = Foo(1, 2)

def main() -> None:
    foos: Foo[1] = empty(1, dtype=Foo)
    init(foos)
    print("foos[0].x =", foos[0].x)

main()

$ lpython --show-c main.py > main.c

$ clang -o app main.c

$ ./app
foos[0].x = 1

$ lpython main.py
foos[0].x = 1

$ python main.py
Traceback (most recent call last):
  File "main.py", line 10, in <module>
    def init(foos: Foo[:]) -> None:
TypeError: 'type' object is not subscriptable

// file: main.c
#include <complex.h>
#include <inttypes.h>

#include <stdlib.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <lfortran_intrinsics.h>

#define ASSERT(cond)                                                           \
    {                                                                          \
        if (!(cond)) {                                                         \
            printf("%s%s", "ASSERT failed: ", __FILE__);                       \
            printf("%s%s", "\nfunction ", __func__);                           \
            printf("%s%d%s", "(), line number ", __LINE__, " at \n");          \
            printf("%s%s", #cond, "\n");                                       \
            exit(1);                                                           \
        }                                                                      \
    }
#define ASSERT_MSG(cond, msg)                                                  \
    {                                                                          \
        if (!(cond)) {                                                         \
            printf("%s%s", "ASSERT failed: ", __FILE__);                       \
            printf("%s%s", "\nfunction ", __func__);                           \
            printf("%s%d%s", "(), line number ", __LINE__, " at \n");          \
            printf("%s%s", #cond, "\n");                                       \
            printf("%s", "ERROR MESSAGE:\n");                                  \
            printf("%s%s", msg, "\n");                                         \
            exit(1);                                                           \
        }                                                                      \
    }

struct dimension_descriptor
{
    int32_t lower_bound, length;
};
struct Foo {
 int32_t x;
 int32_t y;
};

struct xFoo
{
    struct Foo *data;
    struct dimension_descriptor dims[32];
    int32_t n_dims;
    bool is_allocated;
};

// Implementations
void init(void* foos_ptr)
{
    struct xFoo foos_value;
    struct xFoo* foos = &foos_value;
    struct Foo foos_data[1];
    foos->data = foos_data;
    foos->n_dims = 1;
    foos->dims[0].lower_bound = 0;
    foos->dims[0].length = 1;
    foos->n_dims = 1;
    foos->dims[0].lower_bound = 0;
    foos->dims[0].length = 1;
    foos->data = (struct Foo*) foos_ptr;
    foos->data[(0 - foos->dims[0].lower_bound)].x = 1;
    foos->data[(0 - foos->dims[0].lower_bound)].y = 2;
}

void _xx_lcompilers_changed_main_xx()
{
    struct xFoo foos_value;
    struct xFoo* foos = &foos_value;
    struct Foo foos_data[1];
    foos->data = foos_data;
    foos->n_dims = 1;
    foos->dims[0].lower_bound = 0;
    foos->dims[0].length = 1;
    void* foos_ptr;
    foos_ptr = (void*) foos->data;
    init(foos_ptr);
    printf("%s%s%d\n", "foos[0].x =", " ", foos->data[(0 - foos->dims[0].lower_bound)].x);
}

void _lpython_main_program()
{
    _xx_lcompilers_changed_main_xx();
}

float _lfortran_caimag(float complex x);

double _lfortran_zaimag(double complex x);

float pi_32 =   3.14159265358979312e+00;
double pi_64 =   3.14159265358979312e+00;
double _lfortran_dacos(double x);

double _lfortran_dacosh(double x);

double _lfortran_dasin(double x);

double _lfortran_dasinh(double x);

double _lfortran_datan(double x);

double _lfortran_datanh(double x);

double _lfortran_dcos(double x);

double _lfortran_dcosh(double x);

double _lfortran_dexp(double x);

double _lfortran_dlog(double x);

double _lfortran_dlog10(double x);

double _lfortran_dsin(double x);

double _lfortran_dsinh(double x);

double _lfortran_dtan(double x);

double _lfortran_dtanh(double x);

float _lfortran_sacos(float x);

float _lfortran_sacosh(float x);

float _lfortran_sasin(float x);

float _lfortran_sasinh(float x);

float _lfortran_satan(float x);

float _lfortran_satanh(float x);

float _lfortran_scos(float x);

float _lfortran_scosh(float x);

float _lfortran_sexp(float x);

float _lfortran_slog(float x);

float _lfortran_slog10(float x);

float _lfortran_ssin(float x);

float _lfortran_ssinh(float x);

float _lfortran_stan(float x);

float _lfortran_stanh(float x);

int main(int argc, char* argv[])
{
    _lpython_main_program();
    return 0;
}

[certik]

One option is to treat Foo[:] and Foo[1] as syntactic sugar for Array[Foo, :] and Array[Foo, 1]. The "full form" Array style can be implemented in CPython and can be used with custom user types like Foo, while the syntactic sugar can still be used for builtin types like i32[:].

[dylon]

Something like this works for CPython, and the Array implementation for LPython could be modified similarly. TypeVarTuple is new in CPython 3.11.x, but it could easily be replaced with TypeVar as long as the respective parameter is a tuple.

from typing import Any, Dict, TypeVarTuple, TypeVar, Generic

from lpython import c32, c64, i8, i16, i32, i64, f32, f64, dataclass

import numpy as np
from numpy import empty

lp_to_np_types: Dict[Any, Any] = {
    i8: np.int8,
    i16: np.int16,
    i32: np.int32,
    i64: np.int64,
    f32: np.float32,
    f64: np.float64,
    c32: np.complex64,
    c64: np.complex128,
}

def lp_to_np_type(lp_type: Any) -> Any:
    return lp_to_np_types.get(lp_type, lp_type)

DType = TypeVar('DType')
Shape = TypeVarTuple('Shape')

class Array(Generic[DType, *Shape]):
    def __init__(self, dtype, *shape):
        self.dtype = dtype  # keep a reference to dtype since numpy will coerce any non-primitive type to "object"
        self.shape = shape
        self.data = empty(shape=shape, dtype=lp_to_np_type(dtype))
    def __getitem__(self, idxs):
        return self.data.__getitem__(idxs)
    def __setitem__(self, idxs, vals):
        self.data.__setitem__(idxs, vals)

@dataclass
class Foo:
     x: i32
     y: i32

def init(foos: Array[Foo, :]) -> None:
    foos[0] = Foo(1, 2)

def main() -> None:
    foos: Array[Foo, 1] = Array(Foo, 1)
    init(foos)
    print("foos[0].x =", foos[0].x)

main()

[czgdp1807]

So the code in https://github.com/lcompilers/lpython/issues/1799#issuecomment-1546366342 should be made to work with LPython? I think making Foo[:] work would be tricky, because dataclass decorator in lpython.py returns nothing but py_dataclass(arg) i.e., CPython's dataclass object. Now, we might need to inherit this object in a new class and then override __getitem__ to return an Array object. So instead of all these hacks, using Array[Foo, :] would be better. @certik @dylon Let me know what do you folks think.

[certik]

As @dylon wrote in https://github.com/lcompilers/lpython/issues/1799#issuecomment-1546366342, we should use Array[Foo, :], both in CPython and LPython. Let's make it work even for Array[i32, :], and then we will still support i32[:] (as syntactic sugar).

[Shaikh-Ubaid]

I implemented the __class_getitem__ for dataclass decorator in https://github.com/lcompilers/lpython/pull/1814, which seems to support the syntax in example shared in https://github.com/lcompilers/lpython/issues/1799#issue-1708014713. Since, it was a one-line change, I submitted it as a PR. Please kindly review it at your convenience, and if it doesn't meet the requirements, feel free to close the pull request.

[Shaikh-Ubaid]

I think this issue is now fixed. We support the Array[type, dims] syntax. The example shared in the comment of this issue can be updated as follows:

(lp) ubaid@ubaids-MacBook-Pro lpython % git diff
diff --git a/examples/expr2.py b/examples/expr2.py
index afe5e8872..735c0c34f 100644
--- a/examples/expr2.py
+++ b/examples/expr2.py
@@ -1,5 +1,5 @@
 # file: main.py
-from lpython import i32, dataclass
+from lpython import i32, dataclass, Array

 from numpy import empty

@@ -8,11 +8,11 @@ class Foo:
      x: i32
      y: i32

-def init(foos: Foo[:]) -> None:
+def init(foos: Array[Foo, :]) -> None:
     foos[0] = Foo(1, 2)

 def main() -> None:
-    foos: Foo[1] = empty(1, dtype=Foo)
+    foos: Array[Foo, 1] = empty(1, dtype=Foo)
     init(foos)
     print("foos[0].x =", foos[0].x)

(END)

Using the updated example, it works with both LPython and CPython.

from lpython import i32, dataclass, Array

from numpy import empty

@dataclass
class Foo:
     x: i32
     y: i32

def init(foos: Array[Foo, :]) -> None:
    foos[0] = Foo(1, 2)

def main() -> None:
    foos: Array[Foo, 1] = empty(1, dtype=Foo)
    init(foos)
    print("foos[0].x =", foos[0].x)

main()

$  python examples/expr2.py 
foos[0].x = 1

$ lpython examples/expr2.py
foos[0].x = 1

[certik]

Thanks @Shaikh-Ubaid for fixing this!