[#653] Code improvement

* Based on code review from @muloem