How to pass a Swift function to a Metal fragment shader?

I'm trying to create heat maps for a variety of functions of two variables. My first implementation didn't use Metal and was far too slow so now I'm looking into doing it with Metal.

I managed to get a very simple example running but I can't figure out how to pass different functions to the fragment shader. Here's the example:

in ContentView.swift:

struct ContentView: View {
   var body: some View {
      Rectangle()
         .aspectRatio(contentMode: .fit)
         .visualEffect { content, gp in
            let width = Shader.Argument.float(gp.size.width)
            let height = Shader.Argument.float(gp.size.height)
            return content.colorEffect(
               ShaderLibrary.heatMap(width, height)
            )
         }
   }
}

in Shader.metal:

#include <metal_stdlib>
using namespace metal;

constant float twoPi = 6.283185005187988;

// input in [0,1], output in [0,1]
float f(float x) { return (sin(twoPi * x) + 1) / 2; }

// inputs in [0,1], output in [0,1]
float g(float x, float y) { return f(x) * f(y); }

[[ stitchable ]] half4 heatMap(float2 pos, half4 color, float width, float height) {
   float u = pos.x / width;
   float v = pos.y / height;
   float c = g(u, v);
   return half4(c/2, 1-c, c, 1);
}

As it is, it works great and is blazing fast...

...but the function I'm heat-mapping is hardcoded in the metal file. I'd like to be able to write different functions in Swift and pass them to the shader from within SwiftUI (ie, from the ContentView, by querying a model to get the function).

I tried something like this in the metal file:

// (u, v) in [0,1] x [0,1]
// w = f(u, v) in [0,1]
[[ stitchable ]] half4 heatMap(
   float2 pos, half4 color,
   float width, float height,
   float (*f) (float u, float v),
   half4 (*c) (float w)
) {
   float u = pos.x / width;
   float v = pos.y / height;
   float w = f(u, v);
   return c(w);
}

but I couldn't get Swift and C++ to work together to make sense of the function pointers and and now I'm stuck. Any help is greatly appreciated.

Many thanks!