eriknw · Jun 1, 2022 · Jun 1, 2022
diff --git a/graphblas/matrix.py b/graphblas/matrix.py
@@ -895,11 +895,20 @@ def select(self, op, thunk=None):
         if isinstance(op, str):
             op = select.from_string(op)
         else:
-            op, opclass = find_opclass(op)
+            if isinstance(op, MatrixExpression):
+                # Try to rewrite e.g. `A.select(A == 7)` to `gb.select.value(A == 7)`
+                if thunk is not None:
+                    raise TypeError(
+                        "thunk argument not None when calling select with mask or boolean object"
+                    )
+                expr = select._match_expr(self, op)
+                if expr is not None:
+                    return expr
+                opclass = UNKNOWN_OPCLASS
+            else:
+                op, opclass = find_opclass(op)
             if opclass == UNKNOWN_OPCLASS:
-                # Optimization opportunity: we could be smarter and change e.g. `A.select(A == 7)`
-                # to `gb.select.value(A == 7)` or `A.select(gb.select.valueeq, 7)`.
-                # Should we allow e.g., `A.select(B.T)`?
+                # e.g., `A.select(B.S)` or `A.select(B < 7)`
                 mask = _check_mask(op)
                 if thunk is not None:
                     raise TypeError(

diff --git a/graphblas/operator.py b/graphblas/operator.py
@@ -3068,6 +3068,23 @@ def get_semiring(monoid, binaryop, name=None):
 semiring.register_new = Semiring.register_new
 semiring.register_anonymous = Semiring.register_anonymous
 
+select._binary_to_select.update(
+    {
+        binary.eq: select.valueeq,
+        binary.ne: select.valuene,
+        binary.le: select.valuele,
+        binary.lt: select.valuelt,
+        binary.ge: select.valuege,
+        binary.gt: select.valuegt,
+        binary.iseq: select.valueeq,
+        binary.isne: select.valuene,
+        binary.isle: select.valuele,
+        binary.islt: select.valuelt,
+        binary.isge: select.valuege,
+        binary.isgt: select.valuegt,
+    }
+)
+
 _str_to_unary = {
     "-": unary.ainv,
     "~": unary.lnot,

diff --git a/graphblas/select/__init__.py b/graphblas/select/__init__.py
@@ -4,6 +4,7 @@
 # boolean and they default to `select` when called, rather than
 # `apply`).
 _delayed = {}
+_binary_to_select = {}
 
 
 def __dir__():
@@ -48,6 +49,25 @@ def _resolve_expr(expr, callname, opname):
     return globals()[method](tensor, thunk)
 
 
+def _match_expr(parent, expr):
+    """Match expressions to rewrite `A.select(A < 5)` into select expression
+
+    The argument must match the parent, so this _won't_ be rewritten: `A.select(B < 5)`
+    """
+    args = expr.args
+    op = expr.op
+    if (
+        expr.method_name == "apply"
+        and len(args) == 2
+        and op.parent in _binary_to_select
+        and args[1]._is_scalar
+        and (args[0] is parent or getattr(args[0], "T", None) is parent)
+        and op.type in _binary_to_select[op.parent]
+    ):
+        selectop = _binary_to_select[op.parent][op.type]
+        return selectop(parent, args[1])
+
+
 def value(expr):
     """
     An advanced select method which allows for easily expressing

diff --git a/graphblas/tests/test_matrix.py b/graphblas/tests/test_matrix.py
@@ -1140,16 +1140,27 @@ def test_select(A):
     assert w7.isequal(Aupper)
     with pytest.raises(TypeError, match="thunk"):
         A.select(select.valueeq, object())
+
+
+@autocompute
+def test_select_bools_and_masks(A):
+    A3 = Matrix.from_values([0, 3, 3, 6], [3, 0, 2, 4], [3, 3, 3, 3], nrows=7, ncols=7)
     # Select with boolean and masks
     w8 = A.select((A == 3).new()).new()
     assert w8.isequal(A3)
+    w8b = A.select(A == 3).new()  # we rewrite!
+    assert w8b.isequal(A3)
+    w8c = A.select(~(A != 3)).new()
+    assert w8c.isequal(A3)
     w9 = A.select(w8.S).new()
     assert w9.isequal(A3)
     w8[0, 1] = 0
     w10 = A.select(w8.V).new()
     assert w10.isequal(A3)
     with pytest.raises(TypeError, match="thunk"):
         A.select(w8.V, 777)
+    with pytest.raises(TypeError, match="thunk"):
+        A.select(A == 3, 777)
     with pytest.raises(TypeError):
         A.select(A[0, :].new().S)
 

diff --git a/graphblas/tests/test_vector.py b/graphblas/tests/test_vector.py
@@ -688,16 +688,27 @@ def test_select(v):
     assert w6.isequal(result)
     with pytest.raises(TypeError, match="thunk"):
         v.select(select.valueeq, object())
+
+
+@autocompute
+def test_select_bools_and_masks(v):
     # Select with boolean and masks
+    result = Vector.from_values([1, 3], [1, 1], size=7)
     w7 = v.select((v == 1).new()).new()
     assert w7.isequal(result)
+    w7b = v.select(v == 1).new()  # we rewrite!
+    assert w7b.isequal(result)
+    w7c = v.select(~(v != 1)).new()
+    assert w7c.isequal(result)
     w8 = v.select(w7.S).new()
     assert w8.isequal(result)
     w7[4] = 0
     w9 = v.select(w7.V).new()
     assert w9.isequal(result)
     with pytest.raises(TypeError, match="thunk"):
         v.select(w7.V, 777)
+    with pytest.raises(TypeError, match="thunk"):
+        v.select(v == 1, 777)
     with pytest.raises(TypeError):
         v.select(v.outer(v).new().S)
     # Make sure we use `replace`

diff --git a/graphblas/vector.py b/graphblas/vector.py
@@ -827,10 +827,20 @@ def select(self, op, thunk=None):
         if isinstance(op, str):
             op = select.from_string(op)
         else:
-            op, opclass = find_opclass(op)
+            if isinstance(op, VectorExpression):
+                # Try to rewrite e.g. `v.select(v == 7)` to `gb.select.value(v == 7)`
+                if thunk is not None:
+                    raise TypeError(
+                        "thunk argument not None when calling select with mask or boolean object"
+                    )
+                expr = select._match_expr(self, op)
+                if expr is not None:
+                    return expr
+                opclass = UNKNOWN_OPCLASS
+            else:
+                op, opclass = find_opclass(op)
             if opclass == UNKNOWN_OPCLASS:
-                # Optimization opportunity: we could be smarter and change e.g. `v.select(v == 7)`
-                # to `gb.select.value(v == 7)` or `v.select(gb.select.valueeq, 7)`.
+                # e.g., `v.select(w.S)` or `v.select(w < 7)`
                 mask = _check_mask(op)
                 if thunk is not None:
                     raise TypeError(