1 files changed, 60 insertions, 7 deletions

diff --git a/willow/tools/willowc/lib/tokenizer.cpp b/willow/tools/willowc/lib/tokenizer.cpp
index 0c1f917..7ad28a6 100644
--- a/willow/tools/willowc/lib/tokenizer.cpp
+++ b/willow/tools/willowc/lib/tokenizer.cpp
@@ -38,11 +38,11 @@ bool Tokenizer::scan_id(bool accept_digits = true) {
 }
 
 Token Tokenizer::scan() {
-  std::size_t start = this->offset;
-
-  while (isspace(peek()))
+  while (is_space(peek()))
     skip();
 
+  std::size_t start = this->offset;
+
   TokenKind k = [&] {
     switch (peek()) {
     case '@':
@@ -63,6 +63,9 @@ Token Tokenizer::scan() {
     case ',':
       skip();
       return TokenKind::Comma;
+    case ':':
+      skip();
+      return TokenKind::Colon;
     case ';':
       skip();
       return TokenKind::Semicolon;
@@ -99,9 +102,8 @@ Token Tokenizer::scan() {
         return TokenKind::Invalid;
 
       skip();
-      char c = eat();
-      while (c != '\0' && c != '\n')
-        c = eat();
+      while (peek() != '\0' && peek() != '\n')
+        skip();
 
       return TokenKind::Comment;
     }
@@ -124,12 +126,20 @@ Token Tokenizer::scan() {
         return TokenKind::Inst;
       }
 
+      skip();
       return TokenKind::Invalid;
     }
     }
   }();
 
-  return Token{start, offset, k};
+  Token t{start, offset, k};
+  if (t.kind == TokenKind::Invalid) {
+    if (t.start == t.end)
+      t.end++;
+    recover();
+  }
+
+  return t;
 }
 
 bool Tokenizer::scan_dec() {
@@ -173,4 +183,47 @@ bool Tokenizer::scan_constant() {
   return false;
 }
 
+void Tokenizer::recover() {
+  auto is_boundary = [&](char c) {
+    switch (c) {
+    case ' ':
+      [[fallthrough]];
+    case '\n':
+      [[fallthrough]];
+    case '\t':
+      [[fallthrough]];
+    case '\\':
+      [[fallthrough]];
+    case ',':
+      [[fallthrough]];
+    case '%':
+      [[fallthrough]];
+    case '@':
+      [[fallthrough]];
+    case ':':
+      [[fallthrough]];
+    case ';':
+      [[fallthrough]];
+    case '(':
+      [[fallthrough]];
+    case ')':
+      [[fallthrough]];
+    case '{':
+      [[fallthrough]];
+    case '}':
+      [[fallthrough]];
+    case '=':
+      [[fallthrough]];
+    case '\0':
+      return true;
+    default:
+      return false;
+    }
+  };
+
+  while (!is_boundary(peek())) {
+    skip();
+  }
+}
+
 } // namespace willowc