Improve code formatting and documentation