[willow][tools]: unit testing for tokenizer

1 files changed, 19 insertions, 0 deletions

diff --git a/willow/tools/willowc/lib/tokenizer.cpp b/willow/tools/willowc/lib/tokenizer.cpp
index 7ad28a6..8ce5ecb 100644
--- a/willow/tools/willowc/lib/tokenizer.cpp
+++ b/willow/tools/willowc/lib/tokenizer.cpp
@@ -1,3 +1,4 @@
+#include <iostream>
 #include <tokenizer.hpp>
 
 namespace willowc {
@@ -84,6 +85,15 @@ Token Tokenizer::scan() {
     case '=':
       skip();
       return TokenKind::Equals;
+    case '*':
+      skip();
+      return TokenKind::Star;
+    case '<':
+      skip();
+      return TokenKind::LTriangle;
+    case '>':
+      skip();
+      return TokenKind::RTriangle;
     case '-': {
       if (peek(1) == '>') {
         skip(2);
@@ -133,12 +143,19 @@ Token Tokenizer::scan() {
   }();
 
   Token t{start, offset, k};
+
   if (t.kind == TokenKind::Invalid) {
     if (t.start == t.end)
       t.end++;
     recover();
   }
 
+  if (t.kind == TokenKind::Inst) {
+    auto lexeme = buf.substr(t.start, t.end - t.start);
+    if (lexeme == "func")
+      t.kind = TokenKind::FuncKW;
+  }
+
   return t;
 }
 
@@ -214,6 +231,8 @@ void Tokenizer::recover() {
       [[fallthrough]];
     case '=':
       [[fallthrough]];
+    case '^':
+      [[fallthrough]];
     case '\0':
       return true;
     default: