Support single-letter local variables

chibi.h

@@ -11,6 +11,7 @@
 
 typedef enum {
   TK_RESERVED, // Keywords or punctuators
+  TK_IDENT,    // Identifiers
   TK_NUM,      // Integer literals
   TK_EOF,      // End-of-file markers
 } TokenKind;
@@ -28,6 +29,7 @@ struct Token {
 void error(char *fmt, ...);
 void error_at(char *loc, char *fmt, ...);
 bool consume(char *op);
+Token *consume_ident(void);
 void expect(char *op);
 long expect_number(void);
 bool at_eof(void);
@@ -49,8 +51,10 @@ typedef enum {
   ND_NE,        // !=
   ND_LT,        // <
   ND_LE,        // <=
+  ND_ASSIGN,    // =
   ND_RETURN,    // "return"
   ND_EXPR_STMT, // Expression statement
+  ND_VAR,       // Variable
   ND_NUM,       // Integer
 } NodeKind;
 
@@ -61,6 +65,7 @@ struct Node {
   Node *next;    // Next node
   Node *lhs;     // Left-hand side
   Node *rhs;     // Right-hand side
+  char name;     // Used if kind == ND_VAR
   long val;      // Used if kind == ND_NUM
 };
 

codegen.c

@@ -1,5 +1,31 @@
 #include "chibi.h"
 
+// Pushes the given node's address to the stack.
+static void gen_addr(Node *node) {
+  if (node->kind == ND_VAR) {
+    int offset = (node->name - 'a' + 1) * 8;
+    printf("  lea rax, [rbp-%d]\n", offset);
+    printf("  push rax\n");
+    return;
+  }
+
+  error("not an lvalue");
+}
+
+static void load(void) {
+  printf("  pop rax\n");
+  printf("  mov rax, [rax]\n");
+  printf("  push rax\n");
+}
+
+static void store(void) {
+  printf("  pop rdi\n");
+  printf("  pop rax\n");
+  printf("  mov [rax], rdi\n");
+  printf("  push rdi\n");
+}
+
+// Generate code for a given node.
 static void gen(Node *node) {
   switch (node->kind) {
   case ND_NUM:
@@ -9,10 +35,19 @@ static void gen(Node *node) {
     gen(node->lhs);
     printf("  add rsp, 8\n");
     return;
+  case ND_VAR:
+    gen_addr(node);
+    load();
+    return;
+  case ND_ASSIGN:
+    gen_addr(node->lhs);
+    gen(node->rhs);
+    store();
+    return;
   case ND_RETURN:
     gen(node->lhs);
     printf("  pop rax\n");
-    printf("  ret\n");
+    printf("  jmp .L.return\n");
     return;
   }
 
@@ -66,8 +101,17 @@ void codegen(Node *node) {
   printf(".global main\n");
   printf("main:\n");
 
+  // Prologue
+  printf("  push rbp\n");
+  printf("  mov rbp, rsp\n");
+  printf("  sub rsp, 208\n");
+
   for (Node *n = node; n; n = n->next)
     gen(n);
 
+  // Epilogue
+  printf(".L.return:\n");
+  printf("  mov rsp, rbp\n");
+  printf("  pop rbp\n");
   printf("  ret\n");
 }

parse.c

@@ -25,8 +25,15 @@ static Node *new_num(long val) {
   return node;
 }
 
+static Node *new_var_node(char name) {
+  Node *node = new_node(ND_VAR);
+  node->name = name;
+  return node;
+}
+
 static Node *stmt(void);
 static Node *expr(void);
+static Node *assign(void);
 static Node *equality(void);
 static Node *relational(void);
 static Node *add(void);
@@ -60,9 +67,17 @@ static Node *stmt(void) {
   return node;
 }
 
-// expr = equality
+// expr = assign
 static Node *expr(void) {
-  return equality();
+  return assign();
+}
+
+// assign = equality ("=" assign)?
+static Node *assign(void) {
+  Node *node = equality();
+  if (consume("="))
+    node = new_binary(ND_ASSIGN, node, assign());
+  return node;
 }
 
 // equality = relational ("==" relational | "!=" relational)*
@@ -135,13 +150,17 @@ static Node *unary(void) {
   return primary();
 }
 
-// primary = "(" expr ")" | num
+// primary = "(" expr ")" | ident | num
 static Node *primary(void) {
   if (consume("(")) {
     Node *node = expr();
     expect(")");
     return node;
   }
 
+  Token *tok = consume_ident();
+  if (tok)
+    return new_var_node(*tok->str);
+
   return new_num(expect_number());
 }

test.sh

@@ -50,4 +50,7 @@ assert 1 'return 1; 2; 3;'
 assert 2 '1; return 2; 3;'
 assert 3 '1; 2; return 3;'
 
+assert 3 'a=3; return a;'
+assert 8 'a=3; z=5; return a+z;'
+
 echo OK

tokenize.c

@@ -35,6 +35,15 @@ bool consume(char *op) {
   return true;
 }
 
+// Consumes the current token if it is an identifier.
+Token *consume_ident(void) {
+  if (token->kind != TK_IDENT)
+    return NULL;
+  Token *t = token;
+  token = token->next;
+  return t;
+}
+
 // Ensure that the current token is `op`.
 void expect(char *op) {
   if (token->kind != TK_RESERVED || strlen(op) != token->len ||
@@ -98,6 +107,12 @@ Token *tokenize(void) {
       continue;
     }
 
+    // Identifier
+    if ('a' <= *p && *p <= 'z') {
+      cur = new_token(TK_IDENT, cur, p++, 1);
+      continue;
+    }
+
     // Multi-letter punctuators
     if (startswith(p, "==") || startswith(p, "!=") ||
         startswith(p, "<=") || startswith(p, ">=")) {