Rewrite regardless of variant

[AzizZayed]

Is it possible to also make the variant a variable in egglog? For example:

(datatype Shape
    (Shape2D i64 i64)
)

(datatype MatrixOp
    (Matrix String Shape)  ; (<name> <nrows> <ncols>)
    (MatrixTranspose MatrixOp Shape)
)

(let ma (Matrix "a" (Shape2D 5 10)))
(let mb (Matrix "b" (Shape2D 10 15)))
(let mc (Matrix "c" (Shape2D 15 2)))

(function shape-of (MatrixOp) Shape)

(rewrite
    (shape-of (?variant ?op ?s))  ; notice here the variant does is also a variable
    ?s
)

I'm sure there is no support for exactly what I wrote above. But is there any construct within egglog that can achieve the same goal: 1 rewrite rule can apply to many variants of a datatype?

[saulshanabrook]

Not exactly, but the solution I would recommend would be to rewrite those variants into different constructors of the same type... Like for this example:

(datatype Shape
    (Shape2D i64 i64)
)

(datatype MatrixOp
    (matrix String)
    (MatrixTranspose MatrixOp)
)

(datatype Matrix
  (apply-op MatrixOp Shape))

(let ma (apply-op (matrix "a") (Shape2D 5 10))))

(function shape-of (Matrix) Shape)

(rewrite
    (shape-of (apple-op m s)) 
    s
)

If you want to keep the alternative construct you have above, you could also do that, but just rewrite it to this form so that you can apply the rewrite uniformly.