Okay, I've attached a new patch with your suggested implementation,
which to the extent I'm able to test works correctly.