After reconsidering, perhaps the existing implementation is good enough, and after looking again, it seems that the mistake I thought that the existing implementation has, actually it doesn't have that mistake.