Refactor code.

+ minor fixes.