Added some more annotations and null checks