parse.y: fix state after ivar/cvar

* parse.y (parse_atmark): return EXPR_END or EXPR_ENDFN, depending
  on the previous state, even incomplete names consistently.

2 files changed, 21 insertions, 2 deletions

diff --git a/parse.y b/parse.y
index c8930d8792..24b8e2ed41 100644
--- a/parse.y
+++ b/parse.y
@@ -8396,6 +8396,7 @@ parse_atmark(struct parser_params *p, const enum lex_state_e last_state)
 	tokadd(p, '@');
 	c = nextc(p);
     }
+    SET_LEX_STATE(IS_lex_state_for(last_state, EXPR_FNAME) ? EXPR_ENDFN : EXPR_END);
     if (c == -1 || !parser_is_identchar(p)) {
 	pushback(p, c);
 	RUBY_SET_YYLLOC(loc);
@@ -8407,6 +8408,7 @@ parse_atmark(struct parser_params *p, const enum lex_state_e last_state)
 	}
 	parser_show_error_line(p, &loc);
 	set_yylval_noname();
+	SET_LEX_STATE(EXPR_END);
 	return result;
     }
     else if (ISDIGIT(c)) {
@@ -8417,7 +8419,7 @@ parse_atmark(struct parser_params *p, const enum lex_state_e last_state)
 	p->lex.pcur = ptr + len;
 	RUBY_SET_YYLLOC(loc);
 	if (result == tIVAR) {
-	    if (IS_lex_state(EXPR_FNAME)) {
+	    if (IS_lex_state_for(last_state, EXPR_FNAME)) {
 		compile_error(p, "`@%c' is not allowed as an instance variable name", c);
 	    }
 	    else if (ptr[0] == '0') {
@@ -8439,7 +8441,6 @@ parse_atmark(struct parser_params *p, const enum lex_state_e last_state)
     }
 
     if (tokadd_ident(p, c)) return 0;
-    SET_LEX_STATE(EXPR_END);
     tokenize_ident(p, last_state);
     return result;
 }
diff --git a/test/ripper/test_lexer.rb b/test/ripper/test_lexer.rb
index 48110857f8..30b85eff89 100644
--- a/test/ripper/test_lexer.rb
+++ b/test/ripper/test_lexer.rb
@@ -95,4 +95,22 @@ class TestRipper::Lexer < Test::Unit::TestCase
     assert_equal "string\#{nil}\n",
       Ripper.slice(%(<<HERE\nstring\#{nil}\nHERE), "heredoc_beg .*? nl $(.*?) heredoc_end", 1)
   end
+
+  def state(name)
+    Ripper::Lexer::State.new(Ripper.const_get(name))
+  end
+
+  def test_state_after_ivar
+    assert_equal [[1,0],:on_ivar,"@a",state(:EXPR_END)], Ripper.lex("@a").last
+    assert_equal [[1,1],:on_ivar,"@a",state(:EXPR_ENDFN)], Ripper.lex(":@a").last
+    assert_equal [[1,0],:on_ivar,"@1",state(:EXPR_END)], Ripper.lex("@1").last
+    assert_equal [[1,1],:on_ivar,"@1",state(:EXPR_ENDFN)], Ripper.lex(":@1").last
+  end
+
+  def test_state_after_cvar
+    assert_equal [[1,0],:on_cvar,"@@a",state(:EXPR_END)], Ripper.lex("@@a").last
+    assert_equal [[1,1],:on_cvar,"@@a",state(:EXPR_ENDFN)], Ripper.lex(":@@a").last
+    assert_equal [[1,0],:on_cvar,"@@1",state(:EXPR_END)], Ripper.lex("@@1").last
+    assert_equal [[1,1],:on_cvar,"@@1",state(:EXPR_ENDFN)], Ripper.lex(":@@1").last
+  end
 end