Implement compile-time lexer mutators.

This should fix #15.

formatters/html/html.go

@@ -208,8 +208,7 @@ func (f *Formatter) styleToCSS(style *chroma.Style) map[chroma.TokenType]string
 	bg := style.Get(chroma.Background)
 	classes := map[chroma.TokenType]string{}
 	// Convert the style.
-	for t := range style.Entries {
-		e := style.Entries[t]
+	for t, e := range style.Entries {
 		if t != chroma.Background {
 			e = e.Sub(bg)
 		}

mutators.go

@@ -11,6 +11,12 @@ type Mutator interface {
 	Mutate(state *LexerState) error
 }
 
+// A LexerMutator is an additional interface that a Mutator can implement
+// to modify the lexer when it is compiled.
+type LexerMutator interface {
+	MutateLexer(lexer *RegexLexer, rule *CompiledRule) error
+}
+
 // A MutatorFunc is a Mutator that mutates the lexer state machine as it is processing.
 type MutatorFunc func(state *LexerState) error
 
@@ -44,25 +50,32 @@ func Include(state string) Rule {
 	}
 }
 
-// Combined creates a new anonymous state from the given states, and pushes that state.
-func Combined(states ...string) MutatorFunc {
-	return func(s *LexerState) error {
-		name := "__combined_" + strings.Join(states, "__")
-		if _, ok := s.Rules[name]; !ok {
-			combined := []CompiledRule{}
-			for _, state := range states {
-				rules, ok := s.Rules[state]
-				if !ok {
-					return fmt.Errorf("invalid combine state %q", state)
-				}
-				combined = append(combined, rules...)
+type combinedMutator struct {
+	states []string
+}
+
+func (c *combinedMutator) Mutate(s *LexerState) error { return nil }
+
+func (c *combinedMutator) MutateLexer(lexer *RegexLexer, rule *CompiledRule) error {
+	name := "__combined_" + strings.Join(c.states, "__")
+	if _, ok := lexer.rules[name]; !ok {
+		combined := []*CompiledRule{}
+		for _, state := range c.states {
+			rules, ok := lexer.rules[state]
+			if !ok {
+				return fmt.Errorf("invalid combine state %q", state)
 			}
-			s.Rules[name] = combined
+			combined = append(combined, rules...)
 		}
-		s.Rules[s.State][s.Rule].Mutator = Push(name)
-		s.Stack = append(s.Stack, name)
-		return nil
+		lexer.rules[name] = combined
 	}
+	rule.Mutator = Push(name)
+	return nil
+}
+
+// Combined creates a new anonymous state from the given states, and pushes that state.
+func Combined(states ...string) Mutator {
+	return &combinedMutator{states}
 }
 
 // Push states onto the stack.

mutators_test.go

@@ -3,21 +3,22 @@ package chroma
 import (
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestInclude(t *testing.T) {
 	include := Include("other")
 	actual := CompiledRules{
 		"root": {
-			CompiledRule{Rule: include},
+			{Rule: include},
 		},
 		"other": {
-			CompiledRule{Rule: Rule{
+			{Rule: Rule{
 				Pattern: "//.+",
 				Type:    Comment,
 			}},
-			CompiledRule{Rule: Rule{
+			{Rule: Rule{
 				Pattern: `"[^"]*"`,
 				Type:    String,
 			}},
@@ -31,25 +32,38 @@ func TestInclude(t *testing.T) {
 	require.NoError(t, err)
 	expected := CompiledRules{
 		"root": {
-			CompiledRule{Rule: Rule{
+			{Rule: Rule{
 				Pattern: "//.+",
 				Type:    Comment,
 			}},
-			CompiledRule{Rule: Rule{
+			{Rule: Rule{
 				Pattern: `"[^"]*"`,
 				Type:    String,
 			}},
 		},
 		"other": {
-			CompiledRule{Rule: Rule{
+			{Rule: Rule{
 				Pattern: "//.+",
 				Type:    Comment,
 			}},
-			CompiledRule{Rule: Rule{
+			{Rule: Rule{
 				Pattern: `"[^"]*"`,
 				Type:    String,
 			}},
 		},
 	}
 	require.Equal(t, expected, actual)
 }
+
+func TestCombine(t *testing.T) {
+	l := MustNewLexer(nil, Rules{
+		"root":  {{`hello`, String, Combined("world", "bye", "space")}},
+		"world": {{`world`, Name, nil}},
+		"bye":   {{`bye`, Name, nil}},
+		"space": {{`\s+`, Whitespace, nil}},
+	})
+	it, err := l.Tokenise(nil, "hello world")
+	require.NoError(t, err)
+	expected := []*Token{{String, `hello`}, {Whitespace, ` `}, {Name, `world`}}
+	assert.Equal(t, expected, it.Tokens())
+}

regexp.go

@@ -108,7 +108,7 @@ func NewLexer(config *Config, rules Rules) (*RegexLexer, error) {
 	if _, ok := rules["root"]; !ok {
 		return nil, fmt.Errorf("no \"root\" state")
 	}
-	compiledRules := map[string][]CompiledRule{}
+	compiledRules := map[string][]*CompiledRule{}
 	for state, rules := range rules {
 		for _, rule := range rules {
 			flags := ""
@@ -121,7 +121,7 @@ func NewLexer(config *Config, rules Rules) (*RegexLexer, error) {
 			if config.DotAll {
 				flags += "s"
 			}
-			compiledRules[state] = append(compiledRules[state], CompiledRule{Rule: rule, flags: flags})
+			compiledRules[state] = append(compiledRules[state], &CompiledRule{Rule: rule, flags: flags})
 		}
 	}
 	return &RegexLexer{
@@ -144,13 +144,13 @@ type CompiledRule struct {
 	flags  string
 }
 
-type CompiledRules map[string][]CompiledRule
+type CompiledRules map[string][]*CompiledRule
 
 type LexerState struct {
 	Lexer *RegexLexer
 	Text  []rune
 	Pos   int
-	Rules map[string][]CompiledRule
+	Rules CompiledRules
 	Stack []string
 	State string
 	Rule  int
@@ -234,7 +234,7 @@ type RegexLexer struct {
 
 	mu       sync.Mutex
 	compiled bool
-	rules    map[string][]CompiledRule
+	rules    map[string][]*CompiledRule
 }
 
 // SetAnalyser sets the analyser function used to perform content inspection.
@@ -269,7 +269,11 @@ func (r *RegexLexer) maybeCompile() (err error) {
 					return fmt.Errorf("failed to compile rule %s.%d: %s", state, i, err)
 				}
 			}
-			rules[i] = rule
+			if compile, ok := rule.Mutator.(LexerMutator); ok {
+				if err := compile.MutateLexer(r, rule); err != nil {
+					return err
+				}
+			}
 		}
 	}
 	r.compiled = true
@@ -293,7 +297,7 @@ func (r *RegexLexer) Tokenise(options *TokeniseOptions, text string) (Iterator,
 	return state.Iterator(), nil
 }
 
-func matchRules(text []rune, rules []CompiledRule) (int, CompiledRule, []string) {
+func matchRules(text []rune, rules []*CompiledRule) (int, *CompiledRule, []string) {
 	for i, rule := range rules {
 		match, err := rule.Regexp.FindRunesMatch(text)
 		if match != nil && err == nil {
@@ -304,5 +308,5 @@ func matchRules(text []rune, rules []CompiledRule) (int, CompiledRule, []string)
 			return i, rule, groups
 		}
 	}
-	return 0, CompiledRule{}, nil
+	return 0, &CompiledRule{}, nil
 }