diff --git a/lib/imbecile/grammar.rb b/lib/imbecile/grammar.rb
index 0a19d75..1a1cda8 100644
--- a/lib/imbecile/grammar.rb
+++ b/lib/imbecile/grammar.rb
@@ -41,8 +41,9 @@ module Imbecile
       # Build NFA from each token expression.
       @tokens.each do |token_name, token_def|
         token_def[:regex] = Regex.new(token_def[:pattern])
+        token_def[:regex].nfa.end_state.accepts = token_name
       end
-      dfa = DFA.new(@tokens.map {|token_name, token_def| token_def[:regex].nfa})
+      dfa = Regex::DFA.new(@tokens.map {|token_name, token_def| token_def[:regex].nfa})
     end
 
   end
diff --git a/lib/imbecile/regex/nfa.rb b/lib/imbecile/regex/nfa.rb
index 317ed37..c115692 100644
--- a/lib/imbecile/regex/nfa.rb
+++ b/lib/imbecile/regex/nfa.rb
@@ -5,6 +5,8 @@ module Imbecile
 
       class State
 
+        attr_accessor :accepts
+
         def initialize
           @transitions = []
         end