[firedrake] Piecewise-defined function