Retokenize when grammar with matching injection selector is updated

spec/app/text-mate-grammar-spec.coffee

@@ -446,16 +446,44 @@ describe "TextMateGrammar", ->
         expect(tokens[6].value).toBe "http://github.com"
         expect(tokens[6].scopes).toEqual ["source.js", "comment.line.double-slash.js", "markup.underline.link.http.hyperlink"]
 
-      it "retokenizes existing buffers that contain tokens that match the injection selector", ->
-        editSession = project.buildEditSession('sample.js')
-        editSession.setText("// http://github.com")
+      describe "when the grammar is added", ->
+        it "retokenizes existing buffers that contain tokens that match the injection selector", ->
+          editSession = project.buildEditSession('sample.js')
+          editSession.setText("// http://github.com")
 
-        {tokens} = editSession.lineForScreenRow(0)
-        expect(tokens[1].value).toBe " http://github.com"
-        expect(tokens[1].scopes).toEqual ["source.js", "comment.line.double-slash.js"]
+          {tokens} = editSession.lineForScreenRow(0)
+          expect(tokens[1].value).toBe " http://github.com"
+          expect(tokens[1].scopes).toEqual ["source.js", "comment.line.double-slash.js"]
 
-        atom.activatePackage('hyperlink-helper.tmbundle', sync: true)
+          atom.activatePackage('hyperlink-helper.tmbundle', sync: true)
 
-        {tokens} = editSession.lineForScreenRow(0)
-        expect(tokens[2].value).toBe "http://github.com"
-        expect(tokens[2].scopes).toEqual ["source.js", "comment.line.double-slash.js", "markup.underline.link.http.hyperlink"]
+          {tokens} = editSession.lineForScreenRow(0)
+          expect(tokens[2].value).toBe "http://github.com"
+          expect(tokens[2].scopes).toEqual ["source.js", "comment.line.double-slash.js", "markup.underline.link.http.hyperlink"]
+
+      describe "when the grammar is updated", ->
+        it "retokenizes existing buffers that contain tokens that match the injection selector", ->
+          editSession = project.buildEditSession('sample.js')
+          editSession.setText("// SELECT * FROM OCTOCATS")
+
+          {tokens} = editSession.lineForScreenRow(0)
+          expect(tokens[1].value).toBe " SELECT * FROM OCTOCATS"
+          expect(tokens[1].scopes).toEqual ["source.js", "comment.line.double-slash.js"]
+
+          syntax.addGrammar(new TextMateGrammar(
+            name: "test"
+            scopeName: "source.test"
+            repository: {}
+            injectionSelector: "comment"
+            patterns: [ { include: "source.sql" } ]
+          ))
+
+          {tokens} = editSession.lineForScreenRow(0)
+          expect(tokens[1].value).toBe " SELECT * FROM OCTOCATS"
+          expect(tokens[1].scopes).toEqual ["source.js", "comment.line.double-slash.js"]
+
+          atom.activatePackage('sql.tmbundle', sync: true)
+
+          {tokens} = editSession.lineForScreenRow(0)
+          expect(tokens[2].value).toBe "SELECT"
+          expect(tokens[2].scopes).toEqual ["source.js", "comment.line.double-slash.js", "keyword.other.DML.sql"]

src/app/syntax.coffee

@@ -44,7 +44,7 @@ class Syntax
 
   grammarUpdated: (scopeName) ->
     for grammar in @grammars when grammar.scopeName isnt scopeName
-      grammar.grammarUpdated(scopeName)
+      @trigger 'grammar-updated', grammar if grammar.grammarUpdated(scopeName)
 
   setGrammarOverrideForPath: (path, scopeName) ->
     @grammarOverridesByPath[path] = scopeName

src/app/text-mate-grammar.coffee

@@ -71,11 +71,12 @@ class TextMateGrammar
   addIncludedGrammarScope: (scope) ->
     @includedGrammarScopes.push(scope) unless _.include(@includedGrammarScopes, scope)
 
-  grammarUpdated: (scopeName) =>
-    return unless _.include(@includedGrammarScopes, scopeName)
+  grammarUpdated: (scopeName) ->
+    return false unless _.include(@includedGrammarScopes, scopeName)
     @clearRules()
     syntax.grammarUpdated(@scopeName)
     @trigger 'grammar-updated'
+    true
 
   getScore: (path, contents) ->
     contents = fsUtils.read(path) if not contents? and fsUtils.isFile(path)

src/app/tokenized-buffer.coffee

@@ -30,6 +30,9 @@ class TokenizedBuffer
     @buffer.on "changed.tokenized-buffer#{@id}", (e) => @handleBufferChange(e)
     @languageMode.on 'grammar-changed', => @resetScreenLines()
     @languageMode.on 'grammar-updated', => @resetScreenLines()
+    @subscribe syntax, 'grammar-updated', (grammar) =>
+      if grammar.injectionSelector? and @hasTokenForSelector(grammar.injectionSelector)
+        @resetScreenLines()
     @subscribe syntax, 'grammar-added', (grammar) =>
       if grammar.injectionSelector? and @hasTokenForSelector(grammar.injectionSelector)
         @resetScreenLines()