1 files changed, 61 insertions, 6 deletions

diff --git a/usr.bin/oasm/lex.c b/usr.bin/oasm/lex.c
index 7b78d54..87ac464 100644
--- a/usr.bin/oasm/lex.c
+++ b/usr.bin/oasm/lex.c
@@ -34,6 +34,7 @@
 #include <oasm/lex.h>
 #include <oasm/log.h>
 
+#define COMMENT '!'
 #define is_num(c) ((c) >= '0' && (c) <= '9')
 
 static char putback = '\0';
@@ -47,6 +48,21 @@ static char putback = '\0';
 #define S_IMN_INC  "inc"
 #define S_IMN_DEC  "dec"
 #define S_IMN_HLT  "hlt"
+#define S_IMN_BR   "br"
+
+/* Instruction length */
+#define OSMX64_INST_LEN 4
+
+/*
+ * Update the state when the caller encounters
+ * a newline.
+ */
+static inline void
+lex_newline(struct oasm_state *state)
+{
+    ++state->line;
+    state->pip += OSMX64_INST_LEN;
+}
 
 /*
  * Returns 0 if a char is counted as a
@@ -61,7 +77,7 @@ lex_skippable(struct oasm_state *state, char c)
     case '\t': return 0;
     case '\r': return 0;
     case '\n':
-        ++state->line;
+        lex_newline(state);
         return 0;
     }
 
@@ -152,6 +168,10 @@ lex_nomstr(struct oasm_state *state, char **res)
             retval = tmp;
             break;
         }
+        if (tmp == ':') {
+            retval = tmp;
+            break;
+        }
         if (tmp == '\n') {
             ++state->line;
             retval = tmp;
@@ -190,6 +210,19 @@ token_arith(char *p)
     return TT_UNKNOWN;
 }
 
+/*
+ * Control flow instructions
+ */
+static tt_t
+token_cfi(char *p)
+{
+    if (strcmp(p, S_IMN_BR) == 0) {
+        return TT_BR;
+    }
+
+    return TT_UNKNOWN;
+}
+
 static tt_t
 token_xreg(char *p)
 {
@@ -254,6 +287,7 @@ lex_tok(struct oasm_state *state, struct oasm_token *ttp)
 {
     char *p = NULL;
     char c = ' ';
+    short in_comment = 0;
     int tmp;
     tt_t tok;
 
@@ -262,18 +296,24 @@ lex_tok(struct oasm_state *state, struct oasm_token *ttp)
     }
 
     /*
-     * Grab characters. If they are skippable,
-     * don't use them.
+     * Grab characters. If they are skippable or
+     * comments, don't use them.
      */
-    while (lex_skippable(state, c) == 0) {
+    while (lex_skippable(state, c) == 0 || in_comment) {
         if ((c = lex_cin(state)) == 0) {
             return -1;
         }
+
+        if (c == COMMENT) {
+            in_comment = 1;
+        } else if (c == '\n') {
+            in_comment = 0;
+        }
     }
 
     switch (c) {
     case '\n':
-        ++state->line;
+        lex_newline(state);
         return 0;
     case '\0':
         return -1;
@@ -284,7 +324,14 @@ lex_tok(struct oasm_state *state, struct oasm_token *ttp)
         ttp->raw = NULL;
 
         lex_putback(c);
-        lex_nomstr(state, &p);
+        c = lex_nomstr(state, &p);
+
+        while (c == ':') {
+            ttp->type = TT_LABEL;
+            ttp->raw = p;
+            state->label_ip = state->pip;
+            return 0;
+        }
 
         /* Arithmetic operation? */
         if ((tok = token_arith(p)) != TT_UNKNOWN) {
@@ -293,6 +340,13 @@ lex_tok(struct oasm_state *state, struct oasm_token *ttp)
             return 0;
         }
 
+        /* Control flow instruction? */
+        if ((tok = token_cfi(p)) != TT_UNKNOWN) {
+            ttp->type = tok;
+            ttp->raw = p;
+            return 0;
+        }
+
         /* Register? */
         if ((tok = token_reg(p)) != TT_UNKNOWN) {
             ttp->is_reg = 1;
@@ -311,6 +365,7 @@ lex_tok(struct oasm_state *state, struct oasm_token *ttp)
             ttp->raw = p;
             return 0;
         }
+
         oasm_err("bad token \"%s\"\n", p);
         lex_try_free(p);
         return -1;