Override Type::walk/replace for custom storage class

I’m exploring using mutable types in my project. However, it seems that only the “key” part of my storage class and not the “mutable” part is included when using the Type::walk/Type::replace methods. My primary question is how to fix this?

Not to create an “XY problem,” but while I was poking around the MLIR code base in search of a solution, I discovered this promising-looking member of AbstractType:

  /// Function to walk the immediate sub-elements of this type.
  WalkImmediateSubElementsFn walkImmediateSubElementsFn;

However, it isn’t immediately obvious to me how to go about customizing this, or even how it is usually set. Can anyone shed some light on either of these questions?