From 48f1f48f096d965f5e432043244c099899b2a33b Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Fri, 9 Oct 2015 02:27:52 +0300
Subject: [PATCH] transforms.artiq_ir_generator: devirtualize method calls.

---
 .../compiler/transforms/artiq_ir_generator.py | 27 ++++++++++++++-----
 lit-test/test/devirtualization/method.py      | 16 +++++++++++
 2 files changed, 37 insertions(+), 6 deletions(-)
 create mode 100644 lit-test/test/devirtualization/method.py

diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py
index 3daadd8e2..1bd1b01e6 100644
--- a/artiq/compiler/transforms/artiq_ir_generator.py
+++ b/artiq/compiler/transforms/artiq_ir_generator.py
@@ -6,9 +6,9 @@ but without too much detail, such as exposing the reference/value
 semantics explicitly.
 """
 
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from pythonparser import algorithm, diagnostic, ast
-from .. import types, builtins, ir
+from .. import types, builtins, asttyped, ir
 
 def _readable_name(insn):
     if isinstance(insn, ir.Constant):
@@ -100,18 +100,27 @@ class ARTIQIRGenerator(algorithm.Visitor):
         self.unwind_target = None
         self.function_map = dict()
         self.variable_map = dict()
-        self.method_map = dict()
+        self.method_map = defaultdict(lambda: [])
 
     def annotate_calls(self, devirtualization):
         for var_node in devirtualization.variable_map:
             callee_node = devirtualization.variable_map[var_node]
             callee      = self.function_map[callee_node]
+
             call_target = self.variable_map[var_node]
             for use in call_target.uses:
                 if isinstance(use, (ir.Call, ir.Invoke)) and \
                         use.target_function() == call_target:
                     use.static_target_function = callee
 
+        for type_and_method in devirtualization.method_map:
+            callee_node = devirtualization.method_map[type_and_method]
+            callee      = self.function_map[callee_node]
+
+            for call in self.method_map[type_and_method]:
+                assert isinstance(call, (ir.Call, ir.Invoke))
+                call.static_target_function = callee
+
     def add_block(self, name=""):
         block = ir.BasicBlock([], name)
         self.current_function.add(block)
@@ -1553,12 +1562,18 @@ class ARTIQIRGenerator(algorithm.Visitor):
             assert None not in args
 
             if self.unwind_target is None:
-                return self.append(ir.Call(func, args))
+                insn = self.append(ir.Call(func, args))
             else:
                 after_invoke = self.add_block()
-                invoke = self.append(ir.Invoke(func, args, after_invoke, self.unwind_target))
+                insn = self.append(ir.Invoke(func, args, after_invoke, self.unwind_target))
                 self.current_block = after_invoke
-                return invoke
+
+            method_key = None
+            if isinstance(node.func, asttyped.AttributeT):
+                attr_node = node.func
+                self.method_map[(attr_node.value.type, attr_node.attr)].append(insn)
+
+            return insn
 
     def visit_QuoteT(self, node):
         return self.append(ir.Quote(node.value, node.type))
diff --git a/lit-test/test/devirtualization/method.py b/lit-test/test/devirtualization/method.py
new file mode 100644
index 000000000..2105487e3
--- /dev/null
+++ b/lit-test/test/devirtualization/method.py
@@ -0,0 +1,16 @@
+# RUN: env ARTIQ_DUMP_IR=1 %python -m artiq.compiler.testbench.embedding +compile %s 2>%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+class foo:
+    @kernel
+    def bar(self):
+        pass
+x = foo()
+
+@kernel
+def entrypoint():
+    # CHECK-L: ; calls testbench.foo.bar
+    x.bar()