Deal with multiple possible callables

com.ibm.wala.cast.python.ml.test/source/com/ibm/wala/cast/python/ml/test/TestTensorflow2Model.java

@@ -1204,6 +1204,41 @@ public void testModelCall4()
     test("tf2_test_model_call4.py", "SequentialModel.__call__", 1, 1, 3);
   }
 
+  /**
+   * Test call string imprecision as described in
+   * https://github.com/wala/WALA/discussions/1417#discussioncomment-10085680. This should fail due
+   * to https://github.com/wala/ML/issues/207.
+   */
+  @Test(expected = java.lang.AssertionError.class)
+  public void testModelCall5()
+      throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {
+    test(
+        new String[] {
+          "proj66/src/tf2_test_model_call5b.py",
+          "proj66/tf2_test_model_call5.py",
+          "proj66/tf2_test_model_call5a.py"
+        },
+        "tf2_test_model_call5.py",
+        "SequentialModel.__call__",
+        "proj66",
+        1,
+        1,
+        3);
+
+    test(
+        new String[] {
+          "proj66/src/tf2_test_model_call5b.py",
+          "proj66/tf2_test_model_call5.py",
+          "proj66/tf2_test_model_call5a.py"
+        },
+        "tf2_test_model_call5a.py",
+        "SequentialModel.__call__",
+        "proj66",
+        1,
+        1,
+        3);
+  }
+
   @Test
   public void testModelAttributes()
       throws ClassHierarchyException, IllegalArgumentException, CancelException, IOException {

com.ibm.wala.cast.python.test/data/proj66/src/__init__.py

@@ -0,0 +1 @@
+

com.ibm.wala.cast.python.test/data/proj66/src/tf2_test_model_call5b.py

@@ -0,0 +1,9 @@
+# Test https://github.com/wala/WALA/discussions/1417#discussioncomment-10085680.
+
+
+def f(m, d):
+    return m.predict(d)
+
+
+def g(m, d):
+    return f(m, d)

com.ibm.wala.cast.python.test/data/proj66/tf2_test_model_call5.py

@@ -0,0 +1,44 @@
+# Test https://github.com/wala/WALA/discussions/1417#discussioncomment-10085680.
+
+import tensorflow as tf
+from src.tf2_test_model_call5b import g
+
+# Create an override model to classify pictures
+
+
+class SequentialModel(tf.keras.Model):
+
+    def __init__(self, **kwargs):
+        super(SequentialModel, self).__init__(**kwargs)
+
+        self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
+
+        # Add a lot of small layers
+        num_layers = 100
+        self.my_layers = [
+            tf.keras.layers.Dense(64, activation="relu") for n in range(num_layers)
+        ]
+
+        self.dropout = tf.keras.layers.Dropout(0.2)
+        self.dense_2 = tf.keras.layers.Dense(10)
+
+    def __call__(self, x):
+        print("Raffi 1")
+        x = self.flatten(x)
+
+        for layer in self.my_layers:
+            x = layer(x)
+
+        x = self.dropout(x)
+        x = self.dense_2(x)
+
+        return x
+
+    def predict(self, x):
+        return self(x)
+
+
+input_data = tf.random.uniform([20, 28, 28])
+
+model = SequentialModel()
+result = g(model, input_data)

com.ibm.wala.cast.python.test/data/proj66/tf2_test_model_call5a.py

@@ -0,0 +1,44 @@
+# Test https://github.com/wala/WALA/discussions/1417#discussioncomment-10085680.
+
+import tensorflow as tf
+from src.tf2_test_model_call5b import g
+
+# Create an override model to classify pictures
+
+
+class SequentialModel(tf.keras.Model):
+
+    def __init__(self, **kwargs):
+        super(SequentialModel, self).__init__(**kwargs)
+
+        self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
+
+        # Add a lot of small layers
+        num_layers = 100
+        self.my_layers = [
+            tf.keras.layers.Dense(64, activation="relu") for n in range(num_layers)
+        ]
+
+        self.dropout = tf.keras.layers.Dropout(0.2)
+        self.dense_2 = tf.keras.layers.Dense(10)
+
+    def __call__(self, x):
+        print("Raffi 2")
+        x = self.flatten(x)
+
+        for layer in self.my_layers:
+            x = layer(x)
+
+        x = self.dropout(x)
+        x = self.dense_2(x)
+
+        return x
+
+    def predict(self, x):
+        return self(x)
+
+
+input_data = tf.random.uniform([20, 28, 28])
+
+model = SequentialModel()
+result = g(model, input_data)

com.ibm.wala.cast.python/source/com/ibm/wala/cast/python/ipa/callgraph/PythonInstanceMethodTrampolineTargetSelector.java

@@ -43,7 +43,9 @@
 import com.ibm.wala.util.collections.HashMapFactory;
 import com.ibm.wala.util.collections.Pair;
 import com.ibm.wala.util.intset.OrdinalSet;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.logging.Logger;
 
 public class PythonInstanceMethodTrampolineTargetSelector<T>
@@ -87,6 +89,8 @@ protected boolean shouldProcess(CGNode caller, CallSiteReference site, IClass re
 
   @Override
   public IMethod getCalleeTarget(CGNode caller, CallSiteReference site, IClass receiver) {
+    // TODO: Callable detection may need to be moved. See https://github.com/wala/ML/issues/207. If
+    // it stays here, we should further document the receiver swapping process.
     if (isCallable(receiver)) {
       LOGGER.fine("Encountered callable.");
 
@@ -223,6 +227,9 @@ private IClass getCallable(CGNode caller, IClassHierarchy cha, PythonInvokeInstr
     PointerKey receiver = pkf.getPointerKeyForLocal(caller, call.getUse(0));
     OrdinalSet<InstanceKey> objs = builder.getPointerAnalysis().getPointsToSet(receiver);
 
+    // The set of potential callables to be returned.
+    Set<IClass> callableSet = new HashSet<>();
+
     for (InstanceKey o : objs) {
       AllocationSiteInNode instanceKey = getAllocationSiteInNode(o);
       if (instanceKey != null) {
@@ -254,10 +261,22 @@ private IClass getCallable(CGNode caller, IClassHierarchy cha, PythonInvokeInstr
             LOGGER.info("Applying callable workaround for https://github.com/wala/ML/issues/118.");
         }
 
-        if (callable != null) return callable;
+        callableSet.add(callable);
       }
     }
 
+    // if there's only one possible option.
+    if (callableSet.size() == 1) {
+      IClass callable = callableSet.iterator().next();
+      assert callable != null : "Callable should be non-null.";
+      return callable;
+    }
+
+    // if we have multiple candidates.
+    if (callableSet.size() > 1)
+      // we cannot accurately select one.
+      LOGGER.warning("Multiple (" + callableSet.size() + ") callable targets found.");
+
     return null;
   }