Yeah you are right - I changed it to output the masks separately instead with an additional height gradient. This way the result is the same - the node setup afterwards gets a bit more complicated though. Now you can still access the masks after the nodegroup to add occlusion or whatever.
This example isn’t very realistic of course, but if you don’t just mix simple diffuse shaders with it, i think it could be helpful.
.blend File