lcompilers / lpython

Python compiler
https://lpython.org/
Other
1.5k stars 157 forks source link

Implement: Function returning arrays #213

Open certik opened 2 years ago

certik commented 2 years ago

In Fortran there are two options (thus both implemented in ASR):

function zeros1(n) result(r)
integer, intent(in) :: n
real :: r(n)
integer :: i
do i = 1, n
    r(i) = 0
end do
end function

and

function zeros2(n) result(r)
integer, intent(in) :: n
real, allocatable :: r(:)
integer :: i
allocate(r(n))
do i = 1, n
    r(i) = 0
end do
end function

They can both be used in the same way:

real :: A(4)
A = zeros1(4)
A = zeros2(4)

We'll eventually add an ASR optimizer that can effectively convert zeros2() into zeros1(), thus removing the allocatable heap allocation in many common cases, thus making the two equivalent.

In Python, the following syntax is allowed:

def zeros2(n: i32) -> f64[:]:
    A: f64[n]
    A = empty(n)
    i: i32
    for i in range(n):
        A[i] = 0.0
    return A

but this one is not:

def zeros1(n: i32) -> f64[n]:
    A: f64[n]
    A = empty(n)
    i: i32
    for i in range(n):
        A[i] = 0.0
    return A

which gives an error:

Traceback (most recent call last):
  File "/Users/ondrej/repos/lpython/integration_tests/test_numpy_02.py", line 6, in <module>
    def zeros(n: i32) -> f64[n]:
NameError: name 'n' is not defined

So we would have to create some alternative syntax for zeros1() style in Python. The zeros2() allocatable style works.

certik commented 2 years ago

It looks like one option for zeros1 is to use a string as the return type (I've first seen it in https://peps.python.org/pep-0647/ as def is_person(val: dict) -> "TypeGuard[Person]":):

def zeros1(n: i32) -> "f64[n]":
    A: f64[n]
    A = empty(n)
    i: i32
    for i in range(n):
        A[i] = 0.0
    return A

Which leaves the type annotation as a string, and our ast -> asr pass can then parse the string as expected:

In [3]: zeros1.__annotations__
Out[3]: {'n': <ltypes.Type at 0x107494880>, 'return': 'f64[n]'}
namannimmo10 commented 2 years ago

but this one is not:

def zeros1(n: i32) -> f64[n]:
    A: f64[n]
    A = empty(n)
    i: i32
    for i in range(n):
        A[i] = 0.0
    return A

which gives an error:

Traceback (most recent call last):
  File "/Users/ondrej/repos/lpython/integration_tests/test_numpy_02.py", line 6, in <module>
    def zeros(n: i32) -> f64[n]:
NameError: name 'n' is not defined

Why are we moving with this syntax in #223?

certik commented 2 years ago

See https://github.com/lcompilers/lpython/wiki/Array-Types-Design, CPython recently (1 month ago!) chose the syntax:

K = TypeVar('K')
N = TypeVar('N')

def matrix_vector_multiply(x: Array[K, N], v: Array[N]) -> Array[K]: ...

So LPython will go with it too. The f64[n] is just an LPython shortcut for the full CPython's Array[f64, n].

certik commented 2 years ago

It didn't occur to me above that you can declare n ahead of time somehow and then you can do f64[n] which is the most natural syntax, but the PEP 646 does that, so it seems that's the most natural solution.