Add API for structure

llvm_cbuilder/__init__.py

@@ -1,2 +1,3 @@
-from .builder import CBuilder
+from .builder import *
 from .executor import CExecutor
+

llvm_cbuilder/builder.py

@@ -25,6 +25,12 @@ def _is_block_terminated(bb):
     instrs = bb.instructions
     return len(instrs) > 0 and instrs[-1].is_terminator
 
+def _is_cstruct(ty):
+    try:
+        return issubclass(ty, CStruct)
+    except TypeError:
+        return False
+
 @contextlib.contextmanager
 def _change_block_temporarily(builder, bb):
     origbb = builder.basic_block
@@ -135,11 +141,21 @@ class CBuilder(object):
         Only allocate in the first block
         '''
         with _change_block_temporarily(self.builder, self.declare_block):
+            is_cstruct = _is_cstruct(ty)
+            if is_cstruct:
+                cstruct = ty
+                ty = ty.llvm_type()
             ptr = self.builder.alloca(ty, name=name)
-            if value is not None:
-                if not isinstance(value, lc.Value):
-                    value = self.constant(ty, value).value
-                self.builder.store(value, ptr)
+        # back to the body
+        if value is not None:
+            if isinstance(value, CValue):
+                value = value.value
+            if not isinstance(value, lc.Value):
+                value = self.constant(ty, value).value
+            self.builder.store(value, ptr)
+        if is_cstruct:
+            return cstruct(self, ptr)
+        else:
             return CVar(self, ptr)
 
     def array(self, ty, count, name=''):
@@ -616,3 +632,15 @@ class CArray(CValue):
         ptr = builder.gep(self.value, [idx])
         return CVar(self.parent, ptr)
 
+class CStruct(CValue):
+    @classmethod
+    def llvm_type(cls):
+        return lc.Type.struct([v for k, v in cls._fields_])
+
+    def __init__(self, parent, ptr):
+        super(CStruct, self).__init__(parent)
+        makeind = lambda x: self.parent.constant(lc.Type.int(), x).value
+        for i, (fd, _) in enumerate(self._fields_):
+            gep = self.parent.builder.gep(ptr, [makeind(0), makeind(i)])
+            setattr(self, fd, CVar(self.parent, gep))
+

llvm_cbuilder/executor.py

@@ -35,14 +35,18 @@ class CExecutor(object):
         else:
             self.engine = mod_or_engine
 
-    def get_ctype_function(self, fn, typeinfo):
-        types = [ MAP_CTYPES[s.strip()] for s in typeinfo.split(',') ]
-        if not types:
-            retty = None
-            argtys = []
+    def get_ctype_function(self, fn, *typeinfo):
+        if len(typeinfo)==1 and isinstance(typeinfo[0], str):
+            types = [ MAP_CTYPES[s.strip()] for s in typeinfo[0].split(',') ]
+            if not types:
+                retty = None
+                argtys = []
+            else:
+                retty = types[0]
+                argtys = types[1:]
         else:
-            retty = types[0]
-            argtys = types[1:]
+            retty = typeinfo[0]
+            argtys = typeinfo[1:]
 
         prototype = ct.CFUNCTYPE(retty, *argtys)
         fnptr = self.engine.get_pointer_to_function(fn)

llvm_cbuilder/shortnames.py

@@ -7,6 +7,7 @@ int = Type.int(32)
 int16 = short
 int32 = int
 int64 = Type.int(64)
+
 float = Type.float()
 double = Type.double()
 
@@ -15,3 +16,14 @@ double = Type.double()
 pointer = Type.pointer
 
 void_p = pointer(char)
+
+# platform dependent
+
+def _determine_pointer_size():
+    from ctypes import sizeof, c_void_p
+    return sizeof(c_void_p) * 8
+
+pointer_size = _determine_pointer_size()
+
+intp = {32: int32, 64: int64}[pointer_size]
+

tests/test_struct.py

@@ -0,0 +1,52 @@
+from llvm.core import *
+from llvm_cbuilder import *
+import llvm_cbuilder.shortnames as C
+import unittest, ctypes
+
+class Vector2D(CStruct):
+    _fields_ = [
+        ('x', C.float),
+        ('y', C.float),
+    ]
+
+class Vector2DCtype(ctypes.Structure):
+    _fields_ = [
+        ('x', ctypes.c_float),
+        ('y', ctypes.c_float),
+    ]
+
+def gen_vector2d_dist(mod):
+    functype = Type.function(C.float, [C.pointer(Vector2D.llvm_type())])
+    func = mod.add_function(functype, 'vector2d_dist')
+
+    cb = CBuilder(func)
+    vec = cb.var(Vector2D, cb.args[0].load())
+    dist = vec.x * vec.x + vec.y * vec.y
+
+    cb.ret(dist)
+    cb.close()
+    return func
+
+
+class TestStruct(unittest.TestCase):
+    def test_vector2d_dist(self):
+        # prepare module
+        mod = Module.new('mod')
+        lfunc = gen_vector2d_dist(mod)
+        mod.verify()
+        # run
+        exe = CExecutor(mod)
+        func = exe.get_ctype_function(lfunc, ctypes.c_float, ctypes.POINTER(Vector2DCtype))
+
+        from random import random
+        pydist = lambda x, y: x * x + y * y
+        for _ in range(100):
+            x, y = random(), random()
+            vec = Vector2DCtype(x=x, y=y)
+            ans = func(ctypes.pointer(vec))
+            gold = pydist(x, y)
+
+            self.assertLess(abs(ans-gold)/gold, 1e-6)
+
+if __name__ == '__main__':
+    unittest.main()