compiler.embedding: implement string concatenation.

Fixes #526.
m-labs · Aug 8, 2016 · 1a518ea · 1a518ea
1 parent 5a2306a
commit 1a518ea
Showing 8 changed files with 85 additions and 13 deletions.
diff --git a/artiq/compiler/builtins.py b/artiq/compiler/builtins.py
@@ -260,7 +260,7 @@ def is_array(typ, elt=None):
         return types.is_mono(typ, "array")
 
 def is_listish(typ, elt=None):
-    return is_list(typ, elt) or is_array(typ, elt)
+    return is_list(typ, elt) or is_array(typ, elt) or (elt is None and is_str(typ))
 
 def is_range(typ, elt=None):
     if elt is not None:
@@ -283,6 +283,10 @@ def is_iterable(typ):
 def get_iterable_elt(typ):
     if is_iterable(typ):
         return typ.find()["elt"].find()
+    elif is_str(typ):
+        return TInt(types.TValue(8))
+    else:
+        assert False
 
 def is_collection(typ):
     typ = typ.find()

diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -478,8 +478,12 @@ def visit_While(self, node):
 
     def iterable_len(self, value, typ=_size_type):
         if builtins.is_listish(value.type):
+            if isinstance(value, ir.Constant):
+                name = None
+            else:
+                name = "{}.len".format(value.name)
             return self.append(ir.Builtin("len", [value], typ,
-                                          name="{}.len".format(value.name)))
+                                          name=name))
         elif builtins.is_range(value.type):
             start  = self.append(ir.GetAttr(value, "start"))
             stop   = self.append(ir.GetAttr(value, "stop"))
@@ -1313,7 +1317,7 @@ def visit_BinOpT(self, node):
                     loc=node.right.loc)
 
             return self.append(ir.Arith(node.op, lhs, rhs))
-        elif isinstance(node.op, ast.Add): # list + list, tuple + tuple
+        elif isinstance(node.op, ast.Add): # list + list, tuple + tuple, str + str
             lhs, rhs = self.visit(node.left), self.visit(node.right)
             if types.is_tuple(node.left.type) and types.is_tuple(node.right.type):
                 elts = []
@@ -1327,6 +1331,10 @@ def visit_BinOpT(self, node):
                 rhs_length = self.iterable_len(rhs)
 
                 result_length = self.append(ir.Arith(ast.Add(loc=None), lhs_length, rhs_length))
+                if builtins.is_str(node.left.type):
+                    result_last   = result_length
+                    result_length = self.append(ir.Arith(ast.Add(loc=None), result_length,
+                                                         ir.Constant(1, self._size_type)))
                 result = self.append(ir.Alloc([result_length], node.type))
 
                 # Copy lhs
@@ -1350,6 +1358,10 @@ def body_gen(index):
                     lambda index: self.append(ir.Compare(ast.Lt(loc=None), index, rhs_length)),
                     body_gen)
 
+                if builtins.is_str(node.left.type):
+                    self.append(ir.SetElem(result, result_last,
+                                           ir.Constant(0, builtins.TInt(types.TValue(8)))))
+
                 return result
             else:
                 assert False

diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py
@@ -394,6 +394,10 @@ def _coerce_binop(self, op, left, right):
                     self._unify(left.type, right.type,
                                 left.loc, right.loc)
                     return left.type, left.type, right.type
+            elif builtins.is_str(left.type) or builtins.is_str(right.type):
+                self._unify(left.type, right.type,
+                            left.loc, right.loc)
+                return left.type, left.type, right.type
             else:
                 return self._coerce_numeric((left, right), lambda typ: (typ, typ, typ))
         elif isinstance(op, ast.Mult):

diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py
@@ -344,6 +344,8 @@ def llbuiltin(self, name):
             llty = ll.FunctionType(llvoid, [self.llty_of_type(builtins.TException())])
         elif name == "__artiq_reraise":
             llty = ll.FunctionType(llvoid, [])
+        elif name == "strlen":
+            llty = ll.FunctionType(lli32, [llptr])
         elif name == "strcmp":
             llty = ll.FunctionType(lli32, [llptr, llptr])
         elif name == "send_rpc":
@@ -610,6 +612,10 @@ def process_Alloc(self, insn):
                                                    name=insn.name)
             else:
                 assert False
+        elif builtins.is_str(insn.type):
+            llsize = self.map(insn.operands[0])
+            llvalue = self.llbuilder.alloca(lli8, size=llsize)
+            return llvalue
         elif builtins.is_listish(insn.type):
             llsize = self.map(insn.operands[0])
             llvalue = ll.Constant(self.llty_of_type(insn.type), ll.Undefined)
@@ -818,18 +824,27 @@ def process_SetAttr(self, insn):
             return self.llbuilder.store(llvalue, llptr)
 
     def process_GetElem(self, insn):
-        llelts = self.llbuilder.extract_value(self.map(insn.list()), 1)
-        llelt = self.llbuilder.gep(llelts, [self.map(insn.index())],
-                                   inbounds=True)
-        llvalue = self.llbuilder.load(llelt)
+        lst, idx = insn.list(), insn.index()
+        lllst, llidx = map(self.map, (lst, idx))
+        if builtins.is_str(lst.type):
+            llelt = self.llbuilder.gep(lllst, [llidx], inbounds=True)
+            llvalue = self.llbuilder.load(llelt)
+        else:
+            llelts = self.llbuilder.extract_value(lllst, 1)
+            llelt = self.llbuilder.gep(llelts, [llidx], inbounds=True)
+            llvalue = self.llbuilder.load(llelt)
         if isinstance(llvalue.type, ll.PointerType):
             self.mark_dereferenceable(llvalue)
         return llvalue
 
     def process_SetElem(self, insn):
-        llelts = self.llbuilder.extract_value(self.map(insn.list()), 1)
-        llelt = self.llbuilder.gep(llelts, [self.map(insn.index())],
-                                   inbounds=True)
+        lst, idx = insn.list(), insn.index()
+        lllst, llidx = map(self.map, (lst, idx))
+        if builtins.is_str(lst.type):
+            llelt = self.llbuilder.gep(lllst, [llidx], inbounds=True)
+        else:
+            llelts = self.llbuilder.extract_value(lllst, 1)
+            llelt = self.llbuilder.gep(llelts, [llidx], inbounds=True)
         return self.llbuilder.store(self.map(insn.value()), llelt)
 
     def process_Coerce(self, insn):
@@ -1045,8 +1060,11 @@ def get_outer(llenv, env_ty):
             env, = insn.operands
             return get_outer(self.map(env), env.type)
         elif insn.op == "len":
-            lst, = insn.operands
-            return self.llbuilder.extract_value(self.map(lst), 0)
+            collection, = insn.operands
+            if builtins.is_str(collection.type):
+                return self.llbuilder.call(self.llbuiltin("strlen"), [self.map(collection)])
+            else:
+                return self.llbuilder.extract_value(self.map(collection), 0)
         elif insn.op in ("printf", "rtio_log"):
             # We only get integers, floats, pointers and strings here.
             llargs = map(self.map, insn.operands)

diff --git a/artiq/compiler/validators/escape.py b/artiq/compiler/validators/escape.py
@@ -351,6 +351,6 @@ def visit_Return(self, node):
                 {"type": types.TypePrinter().name(node.value.type)},
                 node.value.loc)
             diag = diagnostic.Diagnostic("error",
-                "cannot return a mutable value that does not live forever", {},
+                "cannot return an allocated value that does not live forever", {},
                 node.value.loc, notes=self._diagnostics_for(region, node.value.loc) + [note])
             self.engine.process(diag)
diff --git a/artiq/test/lit/escape/const_string.py b/artiq/test/lit/escape/const_string.py
@@ -0,0 +1,11 @@
+# RUN: %python -m artiq.compiler.testbench.embedding %s
+
+from artiq.experiment import *
+
+@kernel
+def foo():
+    return "x"
+
+@kernel
+def entrypoint():
+    foo()
diff --git a/artiq/test/lit/escape/error_string.py b/artiq/test/lit/escape/error_string.py
@@ -0,0 +1,19 @@
+# RUN: %python -m artiq.compiler.testbench.embedding +diag %s 2>%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+from artiq.experiment import *
+
+@kernel
+def foo():
+    # CHECK-NOT-L: ${LINE:+1}: error:
+    return "x"
+
+@kernel
+def bar():
+    # CHECK-L: ${LINE:+1}: error: cannot return an allocated value that does not live forever
+    return "x" + "y"
+
+@kernel
+def entrypoint():
+    foo()
+    bar()
diff --git a/artiq/test/lit/integration/str.py b/artiq/test/lit/integration/str.py
@@ -0,0 +1,4 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s
+# RUN: %python %s
+
+assert ("x" + "y") == "xy"