Rewrite block partitioning code to make it easier to understand