Cleaner simulation of rebalance

This replaces the 'if .. then .. else ..' with a large block of code in
'else' branch by separating this code into a separate function.

Signed-off-by: Agata Murawska <agatamurawska@google.com>
Reviewed-by: René Nussbaumer <rn@google.com>