Code cleanup
Fixed missing nil check