From 12d1b9819ce06a38e849bab48640867802d4a63f Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Fri, 25 May 2018 02:02:18 +0000
Subject: [PATCH] compiler: handle direct calls to class methods.

Fixes #1005.
---
 artiq/compiler/transforms/__init__.py         |  2 +-
 .../compiler/transforms/llvm_ir_generator.py  | 13 ++++++++++--
 .../lit/embedding/class_fn_direct_call.py     | 20 +++++++++++++++++++
 3 files changed, 32 insertions(+), 3 deletions(-)
 create mode 100644 artiq/test/lit/embedding/class_fn_direct_call.py

diff --git a/artiq/compiler/transforms/__init__.py b/artiq/compiler/transforms/__init__.py
index ccaab7636..64ee25c46 100644
--- a/artiq/compiler/transforms/__init__.py
+++ b/artiq/compiler/transforms/__init__.py
@@ -6,6 +6,6 @@ from .iodelay_estimator import IODelayEstimator
 from .artiq_ir_generator import ARTIQIRGenerator
 from .dead_code_eliminator import DeadCodeEliminator
 from .local_demoter import LocalDemoter
-from .llvm_ir_generator import LLVMIRGenerator
 from .interleaver import Interleaver
 from .typedtree_printer import TypedtreePrinter
+from .llvm_ir_generator import LLVMIRGenerator
diff --git a/artiq/compiler/transforms/llvm_ir_generator.py b/artiq/compiler/transforms/llvm_ir_generator.py
index 6dc76c592..b2fea1852 100644
--- a/artiq/compiler/transforms/llvm_ir_generator.py
+++ b/artiq/compiler/transforms/llvm_ir_generator.py
@@ -9,6 +9,7 @@ from pythonparser import ast, diagnostic
 from llvmlite_artiq import ir as ll, binding as llvm
 from ...language import core as language_core
 from .. import types, builtins, ir
+from ..embedding import SpecializedFunction
 
 
 llvoid     = ll.VoidType()
@@ -1549,8 +1550,16 @@ class LLVMIRGenerator:
             # RPC and C functions have no runtime representation.
             return ll.Constant(llty, ll.Undefined)
         elif types.is_function(typ):
-            return self.get_function_with_undef_env(typ.find(),
-                                                    self.embedding_map.retrieve_function(value))
+            try:
+                func = self.embedding_map.retrieve_function(value)
+            except KeyError:
+                # If a class function was embedded directly (e.g. by a `C.f(...)` call),
+                # but it also appears in a class hierarchy, we might need to fall back
+                # to the non-specialized one, since direct invocations do not cause
+                # monomorphization.
+                assert isinstance(value, SpecializedFunction)
+                func = self.embedding_map.retrieve_function(value.host_function)
+            return self.get_function_with_undef_env(typ.find(), func)
         elif types.is_method(typ):
             llclosure = self._quote(value.__func__, types.get_method_function(typ),
                                     lambda: path() + ['__func__'])
diff --git a/artiq/test/lit/embedding/class_fn_direct_call.py b/artiq/test/lit/embedding/class_fn_direct_call.py
new file mode 100644
index 000000000..91bdac519
--- /dev/null
+++ b/artiq/test/lit/embedding/class_fn_direct_call.py
@@ -0,0 +1,20 @@
+# RUN: %python -m artiq.compiler.testbench.embedding %s
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+class C:
+    @kernel
+    def f(self):
+        pass
+
+class D(C):
+    @kernel
+    def f(self):
+        # super().f()  # super() not bound
+        C.f(self)  # KeyError in compile
+
+di = D()
+@kernel
+def entrypoint():
+    di.f()