Python: Use local flow in Pythagorean.ql

tausbn · web-flow · commit d4fc096ea860 · 2022-04-25T14:35:37.000Z
The hand-rolled notion of flow was causing some severe performance
issues (on a few databases):

```
Tuple counts for Pythagorean::square#168e234a#f#loop_invariant_prefix/2@c86989kr after 6m35s:
175000     ~5%     {2} r1 = JOIN SSA::SsaVariable::getDefinition#dispred#f0820431#ff_10#join_rhs WITH Flow::ControlFlowNode::getNode#dispred#f0820431#bf ON FIRST 1 OUTPUT Lhs.1, Rhs.1 'arg0'
174500     ~6%     {2} r2 = JOIN r1 WITH SSA::SsaVariable::getVariable#dispred#f0820431#ff ON FIRST 1 OUTPUT Rhs.1, Lhs.1 'arg0'
1467782500 ~5%     {3} r3 = JOIN r2 WITH AstGenerated::Name_::getVariable#dispred#f0820431#ff_10#join_rhs ON FIRST 1 OUTPUT 3, Rhs.1 'arg1', Lhs.1 'arg0'
1467553000 ~0%     {2} r4 = JOIN r3 WITH py_expr_contexts_12#join_rhs ON FIRST 2 OUTPUT Lhs.2 'arg0', Lhs.1 'arg1'
                    return r4
```

Rewriting it to use the data flow library made all of this go away. 🎉
diff --git a/python/ql/src/Numerics/Pythagorean.ql b/python/ql/src/Numerics/Pythagorean.ql
@@ -10,36 +10,30 @@
  */
 
 import python
+import semmle.python.dataflow.new.DataFlow
+import semmle.python.ApiGraphs
 
-predicate squareOp(BinaryExpr e) {
-  e.getOp() instanceof Pow and e.getRight().(IntegerLiteral).getN() = "2"
-}
-
-predicate squareMul(BinaryExpr e) {
-  e.getOp() instanceof Mult and e.getRight().(Name).getId() = e.getLeft().(Name).getId()
+DataFlow::ExprNode squareOp() {
+  exists(BinaryExpr e | e = result.asExpr() |
+    e.getOp() instanceof Pow and e.getRight().(IntegerLiteral).getN() = "2"
+  )
 }
 
-predicate squareRef(Name e) {
-  e.isUse() and
-  exists(SsaVariable v, Expr s | v.getVariable() = e.getVariable() |
-    s = v.getDefinition().getNode().getParentNode().(AssignStmt).getValue() and
-    square(s)
+DataFlow::ExprNode squareMul() {
+  exists(BinaryExpr e | e = result.asExpr() |
+    e.getOp() instanceof Mult and e.getRight().(Name).getId() = e.getLeft().(Name).getId()
   )
 }
 
-predicate square(Expr e) {
-  squareOp(e)
-  or
-  squareMul(e)
-  or
-  squareRef(e)
-}
+DataFlow::ExprNode square() { result in [squareOp(), squareMul()] }
 
-from Call c, BinaryExpr s
+from DataFlow::CallCfgNode c, BinaryExpr s, DataFlow::ExprNode left, DataFlow::ExprNode right
 where
-  c.getFunc().toString() = "sqrt" and
-  c.getArg(0) = s and
+  c = API::moduleImport("math").getMember("sqrt").getACall() and
+  c.getArg(0).asExpr() = s and
   s.getOp() instanceof Add and
-  square(s.getLeft()) and
-  square(s.getRight())
+  left.asExpr() = s.getLeft() and
+  right.asExpr() = s.getRight() and
+  left.getALocalSource() = square() and
+  right.getALocalSource() = square()
 select c, "Pythagorean calculation with sub-optimal numerics"
