Rewrite block partitioning code to be easier to understand