Skip to content

arrow

Classes:

ArrowVisitorMixin

Bases: VisitorMixinBase

Methods:

visit

visit(node: ArrowInt32ArrayLength) -> None
Source code in src/irx/builders/llvmliteir/visitors/arrow.py
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
@VisitorCore.visit.dispatch  # type: ignore[attr-defined,untyped-decorator]
def visit(self, node: astx.ArrowInt32ArrayLength) -> None:
    """
    title: Visit ArrowInt32ArrayLength nodes.
    parameters:
      node:
        type: astx.ArrowInt32ArrayLength
    """
    builder_new = self.require_runtime_symbol(
        "arrow", "irx_arrow_array_builder_int32_new"
    )
    append_int32 = self.require_runtime_symbol(
        "arrow", "irx_arrow_array_builder_append_int32"
    )
    finish_builder = self.require_runtime_symbol(
        "arrow", "irx_arrow_array_builder_finish"
    )
    array_length = self.require_runtime_symbol(
        "arrow", "irx_arrow_array_length"
    )
    release_array = self.require_runtime_symbol(
        "arrow", "irx_arrow_array_release"
    )

    builder_slot = self._llvm.ir_builder.alloca(
        self._llvm.ARROW_ARRAY_BUILDER_HANDLE_TYPE,
        name="arrow_builder_slot",
    )
    self._llvm.ir_builder.call(builder_new, [builder_slot])
    builder_handle = self._llvm.ir_builder.load(
        builder_slot, "arrow_builder"
    )

    for item in node.values:
        self.visit_child(item)
        value = safe_pop(self.result_stack)
        if value is None:
            raise Exception("Arrow helper expected an integer value")
        if not is_int_type(value.type):
            raise Exception(
                "Arrow helper supports only integer expressions"
            )

        if value.type.width < self._llvm.INT32_TYPE.width:
            value = self._llvm.ir_builder.sext(
                value, self._llvm.INT32_TYPE, "arrow_i32_promote"
            )
        elif value.type.width > self._llvm.INT32_TYPE.width:
            value = self._llvm.ir_builder.trunc(
                value, self._llvm.INT32_TYPE, "arrow_i32_trunc"
            )

        self._llvm.ir_builder.call(append_int32, [builder_handle, value])

    array_slot = self._llvm.ir_builder.alloca(
        self._llvm.ARROW_ARRAY_HANDLE_TYPE,
        name="arrow_array_slot",
    )
    self._llvm.ir_builder.call(
        finish_builder, [builder_handle, array_slot]
    )
    array_handle = self._llvm.ir_builder.load(array_slot, "arrow_array")
    length_i64 = self._llvm.ir_builder.call(
        array_length, [array_handle], "arrow_length"
    )
    self._llvm.ir_builder.call(release_array, [array_handle])

    length_i32 = self._llvm.ir_builder.trunc(
        length_i64, self._llvm.INT32_TYPE, "arrow_length_i32"
    )
    self.result_stack.append(length_i32)