Conditional lexer rules in Antlr4

Question

Given this grammar:

grammar ColonTest;

main        : statement* EOF;
statement   : NUM_LITERAL expression SEMICOLON;
expression  : primary (MULT_OP primary)*;
primary     : WORD+;

NUM_LITERAL : [0-9]+;
SEMICOLON   : ';';
MULT_OP     : '*' | '/'; // | ':';
WORD        : ('a'..'z' | 'A'..'Z')+;
WS          : [ \t\r\n]+ -> skip;
COLON       : ':' -> skip;

and this input:

1 : statement;
2 : splitted statement
  : into two lines;
3 : a * b / c;
4 : a * b : c;

In the 4'th line the second colon is skipped because of the "COLON" lexer rule. But I need this colon, because its a part of the language (lets say it should be a part of the MULT_OP keyword too). How to achieve this?

EDIT 1:

After deleting COLON : ':' -> skip; and inserting:

statement   : NUM_LITERAL ':' expression (':' expression)* SEMICOLON;

the tree looks like this:

The desired tree should look like this:

_{(source: ibin.co)}

EDIT 2:

How about this? Some implicit tokens defined - but for now it works.

grammar MultiLine;

main        : statement* EOF;
statement   : NUM_LITERAL ':' expression SEMICOLON;
expression  : primary ((MULT_OP|':') primary)*;
primary     : WORD+;

NUM_LITERAL : [0-9]+;
SEMICOLON   : ';';
MULT_OP     : '*' | '/';
WORD        : ('a'..'z' | 'A'..'Z')+;
WS2         : [\r\n]+ [ \t]+ ':' -> skip; // removes all not needed colons
WS          : [ \t\r\n]+ -> skip;

Can I improve that code in any other way?

Answer 1

Tokens such as : should likely not be skipped. Then, alter your statement rule to be:

statement : NUM_LITERAL expression (':' expression)* SEMICOLON;

Answer 2

You need a decision: either you keep the colon or you skip it. You cannot have both opposite actions at the same time. But the question is why you want to skip the colon in the first place? It doesn't hurt to have that in the parse tree, even if you don't need it after the initial number. When you later walk your parse tree you can always just ignore that token.