Refactoring after code review

Fix issues from code review:

Edited by Artyom Kartasov