From 8f5ff07dc08d258209adf60ed8e796efa7b7a476 Mon Sep 17 00:00:00 2001
From: Rui Ueyama <ruiu@cs.stanford.edu>
Date: Sat, 19 Sep 2020 20:53:13 +0900
Subject: [PATCH] Add -include option

---
 chibicc.h      |  2 ++
 main.c         | 51 +++++++++++++++++++++++++++++++++++++++++++++-----
 preprocess.c   |  2 +-
 test/driver.sh |  7 +++++++
 4 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/chibicc.h b/chibicc.h
index e757162..83f81fb 100644
--- a/chibicc.h
+++ b/chibicc.h
@@ -112,6 +112,8 @@ Token *tokenize_file(char *filename);
 // preprocess.c
 //
 
+char *search_include_paths(char *filename);
+bool file_exists(char *path);
 void init_macros(void);
 void define_macro(char *name, char *buf);
 void undef_macro(char *name);
diff --git a/main.c b/main.c
index b2637a8..27bcadf 100644
--- a/main.c
+++ b/main.c
@@ -3,6 +3,7 @@
 StringArray include_paths;
 bool opt_fcommon = true;
 
+static StringArray opt_include;
 static bool opt_E;
 static bool opt_S;
 static bool opt_c;
@@ -22,7 +23,7 @@ static void usage(int status) {
 }
 
 static bool take_arg(char *arg) {
-  char *x[] = {"-o", "-I", "-idirafter"};
+  char *x[] = {"-o", "-I", "-idirafter", "-include"};
 
   for (int i = 0; i < sizeof(x) / sizeof(*x); i++)
     if (!strcmp(arg, x[i]))
@@ -133,6 +134,11 @@ static void parse_args(int argc, char **argv) {
       continue;
     }
 
+    if (!strcmp(argv[i], "-include")) {
+      strarray_push(&opt_include, argv[++i]);
+      continue;
+    }
+
     if (!strcmp(argv[i], "-cc1-input")) {
       base_file = argv[++i];
       continue;
@@ -274,12 +280,47 @@ static void print_tokens(Token *tok) {
   fprintf(out, "\n");
 }
 
-static void cc1(void) {
-  // Tokenize and parse.
-  Token *tok = tokenize_file(base_file);
+static Token *must_tokenize_file(char *path) {
+  Token *tok = tokenize_file(path);
   if (!tok)
-    error("%s: %s", base_file, strerror(errno));
+    error("%s: %s", path, strerror(errno));
+  return tok;
+}
 
+static Token *append_tokens(Token *tok1, Token *tok2) {
+  if (!tok1 || tok1->kind == TK_EOF)
+    return tok2;
+
+  Token *t = tok1;
+  while (t->next->kind != TK_EOF)
+    t = t->next;
+  t->next = tok2;
+  return tok1;
+}
+
+static void cc1(void) {
+  Token *tok = NULL;
+
+  // Process -include option
+  for (int i = 0; i < opt_include.len; i++) {
+    char *incl = opt_include.data[i];
+
+    char *path;
+    if (file_exists(incl)) {
+      path = incl;
+    } else {
+      path = search_include_paths(incl);
+      if (!path)
+        error("-include: %s: %s", incl, strerror(errno));
+    }
+
+    Token *tok2 = must_tokenize_file(path);
+    tok = append_tokens(tok, tok2);
+  }
+
+  // Tokenize and parse.
+  Token *tok2 = must_tokenize_file(base_file);
+  tok = append_tokens(tok, tok2);
   tok = preprocess(tok);
 
   // If -E is given, print out preprocessed C code as a result.
diff --git a/preprocess.c b/preprocess.c
index 90c7ac7..e0d3281 100644
--- a/preprocess.c
+++ b/preprocess.c
@@ -687,7 +687,7 @@ static bool expand_macro(Token **rest, Token *tok) {
   return true;
 }
 
-static char *search_include_paths(char *filename) {
+char *search_include_paths(char *filename) {
   if (filename[0] == '/')
     return filename;
 
diff --git a/test/driver.sh b/test/driver.sh
index a223dc5..2fdf08a 100755
--- a/test/driver.sh
+++ b/test/driver.sh
@@ -181,4 +181,11 @@ check '-fcommon'
 echo 'int foo;' | $chibicc -fno-common -S -o- - | grep -q '^foo:'
 check '-fno-common'
 
+# -include
+echo foo > $tmp/out.h
+echo bar | $chibicc -include $tmp/out.h -E -o- - | grep -q -z 'foo.*bar'
+check -include
+echo NULL | $chibicc -Iinclude -include stdio.h -E -o- - | grep -q 0
+check -include
+
 echo OK