rebase: don't update state dict same way for each root The update statement does not depend on anything in the loop, so just move it before the loop and do it once. There are no cases where update would happen 0 times before (and 1 now); the function returns early in all such cases.