1 files changed, 9 insertions, 18 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java b/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java
index b24e188f9c..f4b55ad4c4 100644
--- a/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java
+++ b/src/main/java/com/google/devtools/build/lib/syntax/Runtime.java
@@ -26,7 +26,6 @@ import net.bytebuddy.implementation.bytecode.StackManipulation;
 
 import java.lang.reflect.Field;
 import java.util.HashMap;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -120,20 +119,15 @@ public final class Runtime {
    */
   public static void registerFunction(Class<?> nameSpace, BaseFunction function) {
     Preconditions.checkNotNull(nameSpace);
-    // TODO(bazel-team): fix our code so that the two checks below work.
-    // Preconditions.checkArgument(function.getObjectType().equals(nameSpace));
-    // Preconditions.checkArgument(nameSpace.equals(getCanonicalRepresentation(nameSpace)));
-    nameSpace = getCanonicalRepresentation(nameSpace);
+    Preconditions.checkArgument(nameSpace.equals(getCanonicalRepresentation(nameSpace)));
+    Preconditions.checkArgument(
+        getCanonicalRepresentation(function.getObjectType()).equals(nameSpace));
     if (!functions.containsKey(nameSpace)) {
       functions.put(nameSpace, new HashMap<String, BaseFunction>());
     }
     functions.get(nameSpace).put(function.getName(), function);
   }
 
-  static Map<String, BaseFunction> getNamespaceFunctions(Class<?> nameSpace) {
-    return functions.get(getCanonicalRepresentation(nameSpace));
-  }
-
   /**
    * Returns the canonical representation of the given class, i.e. the super class for which any
    * functions were registered.
@@ -143,18 +137,15 @@ public final class Runtime {
    */
   // TODO(bazel-team): make everything a SkylarkValue, and remove this function.
   public static Class<?> getCanonicalRepresentation(Class<?> clazz) {
-    if (SkylarkValue.class.isAssignableFrom(clazz)) {
-      return clazz;
-    }
-    if (Map.class.isAssignableFrom(clazz)) {
-      return MethodLibrary.DictModule.class;
-    }
     if (String.class.isAssignableFrom(clazz)) {
       return MethodLibrary.StringModule.class;
     }
-    Preconditions.checkArgument(
-        !List.class.isAssignableFrom(clazz), "invalid non-SkylarkList list class");
-    return clazz;
+    return EvalUtils.getSkylarkType(clazz);
+  }
+
+
+  static Map<String, BaseFunction> getNamespaceFunctions(Class<?> nameSpace) {
+    return functions.get(getCanonicalRepresentation(nameSpace));
   }
 
   /**