Rewrite block partitioning code to make it easier to understand

Also try to follow the coding style more closely.