Slight optimization to shader code, added todo