Refactor flattening/wrapping code, add flattening/wrapping unit tests.