Fix some bugs in task 4.6

Signed-off-by: Serg Brych <sbrych@gmail.com>