Clean up code and fix.
2 files changed