Some more code improvements

Some more code improvements