Add stack allocation handles

Create handles for simple stack allocation and allocation of an array.

dispyatcher/allocation.py

@@ -0,0 +1,93 @@
+import itertools
+import llvmlite.ir
+from llvmlite.ir import Value as IRValue
+from typing import Generic, Tuple, Sequence, Union
+
+from dispyatcher import F, Handle, Type, ReturnManagement, ArgumentManagement, TemporaryValue, Pointer
+from dispyatcher.general import UncheckedArray
+
+
+class ToStackPointer(Handle[F], Generic[F]):
+    """
+    Creates a stack allocation for a value and provides a pointer to it.
+
+    The allocation can take ownership of the value and the value will be freed when the pointer is dropped.
+    """
+    __type: Type
+    __transfer: ReturnManagement
+
+    def __init__(self, ty: Type, transfer: ReturnManagement = ReturnManagement.TRANSFER):
+        """
+        Create a new stack allocation handle for a particular type.
+
+        :param ty: the type the allocation contains; the result will be a pointer to this type
+        :param transfer: whether to move the value into the allocation (``TRANSFER``) or copy it and rely on the
+            original value for memory management (``BORROW``).
+        """
+        super().__init__()
+        self.__type = ty
+        self.__transfer = transfer
+
+    def __str__(self) -> str:
+        return f"ToStackPointer[{self.__transfer.name}] → {self.__type}"
+
+    def generate_handle_ir(self, flow: F, args: Sequence[IRValue]) -> Union[TemporaryValue, IRValue]:
+        ty = self.__type.machine_type()
+        (arg, ) = args
+        alloc = flow.builder.alloca(ty)
+        flow.builder.store(arg, alloc)
+        return alloc
+
+    def handle_arguments(self) -> Sequence[Tuple[Type, ArgumentManagement]]:
+        return (self.__type,
+                ArgumentManagement.TRANSFER_CAPTURE_PARENTS
+                if self.__transfer == ReturnManagement.TRANSFER
+                else ArgumentManagement.BORROW_CAPTURE_PARENTS),
+
+    def handle_return(self) -> Tuple[Type, ReturnManagement]:
+        return Pointer(self.__type), self.__transfer
+
+
+class CollectIntoArray(Handle[F], Generic[F]):
+    """
+    Stores a arguments into a stack-allocated fixed-length array.
+
+    This is meant for producing arrays like ``argv`` from individual values
+    """
+    __type: Type
+    __count: int
+    __null_terminated: bool
+
+    def __init__(self, ty: Type, count: int, null_terminated: bool = False):
+        """
+        Creates a new stack-allocated array collector handle.
+
+        :param ty: the type of the elements in the array
+        :param count: the number of items to collect
+        :param null_terminated: if true, the array will be allocated with an extra slot at the end that will have a null
+            or zero value; otherwise, the array will only be filled with the arguments.
+        """
+        super().__init__()
+        assert count > 0, "Array length must be positive number"
+        self.__type = ty
+        self.__count = count
+        self.__null_terminated = null_terminated
+
+    def __str__(self) -> str:
+        return f"CollectIntoArray({self.__type} * {self.__count} {'+null' if self.__null_terminated else ''})"
+
+    def generate_handle_ir(self, flow: F, args: Sequence[IRValue]) -> Union[TemporaryValue, IRValue]:
+        i32 = llvmlite.ir.IntType(32)
+        array = flow.builder.alloca(llvmlite.ir.ArrayType(self.__type.machine_type(),
+                                                          self.__count + self.__null_terminated))
+        for idx, arg in enumerate(args):
+            flow.builder.store(arg, flow.builder.gep(array, [i32(0), i32(idx)]))
+        if self.__null_terminated:
+            flow.builder.store(self.__type.machine_type()(None), flow.builder.gep(array, [i32(0), i32(self.__count)]))
+        return flow.builder.bitcast(array, self.__type.machine_type().as_pointer())
+
+    def handle_arguments(self) -> Sequence[Tuple[Type, ArgumentManagement]]:
+        return list(itertools.repeat((self.__type, ArgumentManagement.BORROW_CAPTURE), self.__count))
+
+    def handle_return(self) -> Tuple[Type, ReturnManagement]:
+        return UncheckedArray(self.__type), ReturnManagement.BORROW

dispyatcher/cpython.py

@@ -9,12 +9,13 @@
 from dispyatcher import ArgumentManagement, BaseTransferUnaryHandle, ControlFlow, DerefPointer, F, FlowState, Handle, \
     Identity, Type, ReturnManagement, TemporaryValue
 from dispyatcher.accessors import GetElementPointer
-from dispyatcher.general import BaseIndirectFunction, CurrentProcessFunction, MachineType
+from dispyatcher.general import BaseIndirectFunction, CurrentProcessFunction, MachineType, UncheckedArray
 from dispyatcher.permute import implode_args
 from dispyatcher.repacking import Repacker, RepackingDispatcher, RepackingState
 
 INT_RESULT_TYPE = llvmlite.ir.types.IntType(ctypes.sizeof(ctypes.c_int) * 8)
 SIZE_T_TYPE = llvmlite.ir.types.IntType(ctypes.sizeof(ctypes.c_size_t) * 8)
+CHAR_ARRAY = UncheckedArray(MachineType(llvmlite.ir.IntType(8)))
 
 
 class PythonControlFlow(dispyatcher.ControlFlow):
@@ -804,3 +805,7 @@ def handle_return(self) -> Tuple[Type, ReturnManagement]:
                                       ReturnManagement.TRANSFER,
                                       "PyDict_Size",
                                       (PY_DICT_TYPE, ArgumentManagement.BORROW_TRANSIENT))
+PY_UNICODE_FROM_STRING = CurrentProcessFunction(PyObjectType(str),
+                                                ReturnManagement.TRANSFER,
+                                                "PyUnicode_FromString",
+                                                (CHAR_ARRAY, ArgumentManagement.BORROW_TRANSIENT))

docs/index.rst

@@ -48,6 +48,9 @@ Once a handle is constructed, creating a ``Callsite`` will compile it and turn i
 .. automodule:: dispyatcher.accessors
     :members:
 
+.. automodule:: dispyatcher.allocation
+    :members:
+
 .. automodule:: dispyatcher.cpython
     :members:
 

tests/test_cpython.py

@@ -6,14 +6,15 @@
 import dispyatcher
 import dispyatcher.general
 from dispyatcher import CallSite, Identity, IgnoreArguments
-from dispyatcher.general import SimpleConstant
+from dispyatcher.allocation import CollectIntoArray
+from dispyatcher.general import MachineType, SimpleConstant
 from dispyatcher.cpython import TupleUnpackingDispatcher, PY_OBJECT_TYPE, PythonControlFlowType
 
 
 class RepackTests(unittest.TestCase):
 
     def test_tuple_unpack_int(self):
-        i32 = dispyatcher.general.MachineType(llvmlite.ir.IntType(32))
+        i32 = MachineType(llvmlite.ir.IntType(32))
         tud = TupleUnpackingDispatcher(IgnoreArguments(SimpleConstant(i32, -1), 0, PY_OBJECT_TYPE), 0)
         tud.append(Identity(i32), None)
         callsite = CallSite(tud, PythonControlFlowType())
@@ -24,7 +25,7 @@ def test_tuple_unpack_int(self):
         self.assertEqual(callsite(ctypes.py_object((7.5,))), -1)
 
     def test_tuple_unpack_double(self):
-        dbl = dispyatcher.general.MachineType(llvmlite.ir.DoubleType())
+        dbl = MachineType(llvmlite.ir.DoubleType())
         tud = TupleUnpackingDispatcher(IgnoreArguments(SimpleConstant(dbl, -1), 0, PY_OBJECT_TYPE), 0)
         tud.append(Identity(dbl), None)
         callsite = CallSite(tud, PythonControlFlowType())
@@ -61,3 +62,10 @@ def test_value(self):
         self_handle = dispyatcher.cpython.Value(self, transfer=dispyatcher.ReturnManagement.TRANSFER)
         callsite = CallSite(self_handle, PythonControlFlowType())
         self.assertEqual(self, callsite())
+
+    def test_array_collection(self):
+        # This is really an allocation test, but we use Python infrastructure to make it pleasant
+        i8 = MachineType(llvmlite.ir.IntType(8))
+        handle = CollectIntoArray(i8, 2, null_terminated=True) + dispyatcher.cpython.PY_UNICODE_FROM_STRING
+        callsite = CallSite(handle, PythonControlFlowType())
+        self.assertEqual(callsite(ord('h'), ord('i')), "hi")