Code clean-up of TODO(bajtos)

Resolved all issues mentioned in TODO(bajtos) comments.