[ruby/prism] Unary symbols that cannot be binary should drop @

https://github.com/ruby/prism/commit/d139af033f

prism/prism.c

@@ -12626,14 +12626,48 @@ parse_string_part(pm_parser_t *parser) {
     }
 }
 
+/**
+ * When creating a symbol, unary operators that cannot be binary operators
+ * automatically drop trailing `@` characters. This happens at the parser level,
+ * such that `~@` is parsed as `~` and `!@` is parsed as `!`. We do that here.
+ */
+static pm_node_t *
+parse_operator_symbol(pm_parser_t *parser, const pm_token_t *opening, pm_lex_state_t next_state) {
+    pm_token_t closing = not_provided(parser);
+    pm_symbol_node_t *symbol = pm_symbol_node_create(parser, opening, &parser->current, &closing);
+
+    const uint8_t *end = parser->current.end;
+    switch (parser->current.type) {
+        case PM_TOKEN_TILDE:
+        case PM_TOKEN_BANG:
+            if (parser->current.end[-1] == '@') end--;
+            break;
+        default:
+            break;
+    }
+
+    if (next_state != PM_LEX_STATE_NONE) lex_state_set(parser, next_state);
+    parser_lex(parser);
+
+    pm_string_shared_init(&symbol->unescaped, parser->previous.start, end);
+    return (pm_node_t *) symbol;
+}
+
+/**
+ * Parse a symbol node. This function will get called immediately after finding
+ * a symbol opening token. This handles parsing bare symbols and interpolated
+ * symbols.
+ */
 static pm_node_t *
 parse_symbol(pm_parser_t *parser, pm_lex_mode_t *lex_mode, pm_lex_state_t next_state) {
-    pm_token_t opening = parser->previous;
+    const pm_token_t opening = parser->previous;
 
     if (lex_mode->mode != PM_LEX_STRING) {
         if (next_state != PM_LEX_STATE_NONE) lex_state_set(parser, next_state);
 
         switch (parser->current.type) {
+            case PM_CASE_OPERATOR:
+                return parse_operator_symbol(parser, &opening, next_state == PM_LEX_STATE_NONE ? PM_LEX_STATE_ENDFN : next_state);
             case PM_TOKEN_IDENTIFIER:
             case PM_TOKEN_CONSTANT:
             case PM_TOKEN_INSTANCE_VARIABLE:
@@ -12645,10 +12679,6 @@ parse_symbol(pm_parser_t *parser, pm_lex_mode_t *lex_mode, pm_lex_state_t next_s
             case PM_CASE_KEYWORD:
                 parser_lex(parser);
                 break;
-            case PM_CASE_OPERATOR:
-                lex_state_set(parser, next_state == PM_LEX_STATE_NONE ? PM_LEX_STATE_ENDFN : next_state);
-                parser_lex(parser);
-                break;
             default:
                 expect2(parser, PM_TOKEN_IDENTIFIER, PM_TOKEN_METHOD_NAME, PM_ERR_SYMBOL_INVALID);
                 break;
@@ -12764,8 +12794,11 @@ parse_symbol(pm_parser_t *parser, pm_lex_mode_t *lex_mode, pm_lex_state_t next_s
 static inline pm_node_t *
 parse_undef_argument(pm_parser_t *parser) {
     switch (parser->current.type) {
+        case PM_CASE_OPERATOR: {
+            const pm_token_t opening = not_provided(parser);
+            return parse_operator_symbol(parser, &opening, PM_LEX_STATE_NONE);
+        }
         case PM_CASE_KEYWORD:
-        case PM_CASE_OPERATOR:
         case PM_TOKEN_CONSTANT:
         case PM_TOKEN_IDENTIFIER:
         case PM_TOKEN_METHOD_NAME: {
@@ -12799,16 +12832,17 @@ parse_undef_argument(pm_parser_t *parser) {
 static inline pm_node_t *
 parse_alias_argument(pm_parser_t *parser, bool first) {
     switch (parser->current.type) {
-        case PM_CASE_OPERATOR:
+        case PM_CASE_OPERATOR: {
+            const pm_token_t opening = not_provided(parser);
+            return parse_operator_symbol(parser, &opening, first ? PM_LEX_STATE_FNAME | PM_LEX_STATE_FITEM : PM_LEX_STATE_NONE);
+        }
         case PM_CASE_KEYWORD:
         case PM_TOKEN_CONSTANT:
         case PM_TOKEN_IDENTIFIER:
         case PM_TOKEN_METHOD_NAME: {
-            if (first) {
+            if (first) lex_state_set(parser, PM_LEX_STATE_FNAME | PM_LEX_STATE_FITEM);
-                lex_state_set(parser, PM_LEX_STATE_FNAME | PM_LEX_STATE_FITEM);
-            }
-
             parser_lex(parser);
+
             pm_token_t opening = not_provided(parser);
             pm_token_t closing = not_provided(parser);
             pm_symbol_node_t *symbol = pm_symbol_node_create(parser, &opening, &parser->previous, &closing);

test/prism/snapshots/symbols.txt

@@ -135,7 +135,7 @@
         │   ├── opening_loc: (27,0)-(27,1) = ":"
         │   ├── value_loc: (27,1)-(27,3) = "~@"
         │   ├── closing_loc: ∅
-        │   └── unescaped: "~@"
+        │   └── unescaped: "~"
         ├── @ ArrayNode (location: (29,0)-(29,16))
         │   ├── flags: ∅
         │   ├── elements: (length: 4)
@@ -345,7 +345,7 @@
         │   ├── opening_loc: (57,0)-(57,1) = ":"
         │   ├── value_loc: (57,1)-(57,3) = "!@"
         │   ├── closing_loc: ∅
-        │   └── unescaped: "!@"
+        │   └── unescaped: "!"
         ├── @ SymbolNode (location: (59,0)-(59,3))
         │   ├── flags: ∅
         │   ├── opening_loc: (59,0)-(59,1) = ":"