Improvements to the ANTLR4 grammar