So, it looks definitely more complicated (and fragile?), but I think I still like it better. Maybe you could add a comment in the code documenting the two different paths and the associated checks, so that we don't forget. :)
Sent the v2 with this approach: https://lore.kernel.org/lkml/20250317022325.52791-1-harshit@nutanix.com/T/#u
Please take a look.
Thanks, Harshit