Code drop from //branches/cupcake/...@124589

3 files changed, 106 insertions, 8 deletions

diff --git a/tools/atree/atree.cpp b/tools/atree/atree.cpp
index 4d97d24..aee2b3c 100644
--- a/tools/atree/atree.cpp
+++ b/tools/atree/atree.cpp
@@ -5,12 +5,15 @@
 #include "files.h"
 #include "fs.h"
 #include <set>
+#include <iostream>
+#include <sstream>
 
 using namespace std;
 
 bool g_debug = false;
 vector<string> g_listFiles;
 vector<string> g_inputBases;
+map<string, string> g_variables;
 string g_outputBase;
 string g_dependency;
 bool g_useHardLinks = false;
@@ -29,6 +32,7 @@ const char* USAGE =
 "  -l             Use hard links instead of copying the files.\n"
 "  -m DEPENDENCY  Output a make-formatted file containing the list.\n"
 "                 of files included.  It sets the variable ATREE_FILES.\n"
+"  -v VAR=VAL     Replaces ${VAR} by VAL when reading input files.\n"
 "\n"
 "FILELIST file format:\n"
 "  The FILELIST files contain the list of files that will end up\n"
@@ -53,13 +57,28 @@ int usage()
     return 1;
 }
 
+static bool
+add_variable(const char* arg) {
+    const char* p = arg;
+    while (*p && *p != '=') p++;
+
+    if (*p == 0 || p == arg || p[1] == 0) {
+        return false;
+    }
+
+    ostringstream var;
+    var << "${" << string(arg, p-arg) << "}";
+    g_variables[var.str()] = string(p+1);
+    return true;
+}
+
 int
 main(int argc, char* const* argv)
 {
     int err;
     bool done = false;
     while (!done) {
-        int opt = getopt(argc, argv, "f:I:o:hlm:");
+        int opt = getopt(argc, argv, "f:I:o:hlm:v:");
         switch (opt)
         {
             case -1:
@@ -90,6 +109,14 @@ main(int argc, char* const* argv)
                 }
                 g_dependency = optarg;
                 break;
+            case 'v':
+                if (!add_variable(optarg)) {
+                    fprintf(stderr, "%s Invalid expression in '-v %s': "
+                            "expected format is '-v VAR=VALUE'.\n",
+                            argv[0], optarg);
+                    return usage();
+                }
+                break;
             default:
             case '?':
             case 'h':
@@ -143,7 +170,7 @@ main(int argc, char* const* argv)
     // read file lists
     for (vector<string>::iterator it=g_listFiles.begin();
                                 it!=g_listFiles.end(); it++) {
-        err = read_list_file(*it, &files, &excludes);
+        err = read_list_file(*it, g_variables, &files, &excludes);
         if (err != 0) {
             return err;
         }
diff --git a/tools/atree/files.cpp b/tools/atree/files.cpp
index 5842378..0f59521 100644
--- a/tools/atree/files.cpp
+++ b/tools/atree/files.cpp
@@ -100,9 +100,61 @@ add_file(vector<FileRecord>* files, const string& listFile, int listLine,
     files->push_back(rec);
 }
 
+static string
+replace_variables(const string& input,
+                  const map<string, string>& variables,
+                  bool* error) {
+    if (variables.empty()) {
+        return input;
+    }
+
+    // Abort if the variable prefix is not found
+    if (input.find("${") == string::npos) {
+        return input;
+    }
+
+    string result = input;
+
+    // Note: rather than be fancy to detect recursive replacements,
+    // we simply iterate till a given threshold is met.
+
+    int retries = 1000;
+    bool did_replace;
+
+    do {
+        did_replace = false;
+        for (map<string, string>::const_iterator it = variables.begin();
+             it != variables.end(); ++it) {
+            string::size_type pos = 0;
+            while((pos = result.find(it->first, pos)) != string::npos) {
+                result = result.replace(pos, it->first.length(), it->second);
+                pos += it->second.length();
+                did_replace = true;
+            }
+        }
+        if (did_replace && --retries == 0) {
+            *error = true;
+            fprintf(stderr, "Recursive replacement detected during variables "
+                    "substitution. Full list of variables is: ");
+
+            for (map<string, string>::const_iterator it = variables.begin();
+                 it != variables.end(); ++it) {
+                fprintf(stderr, "  %s=%s\n",
+                        it->first.c_str(), it->second.c_str());
+            }
+
+            return result;
+        }
+    } while (did_replace);
+
+    return result;
+}
+
 int
-read_list_file(const string& filename, vector<FileRecord>* files,
-                    vector<string>* excludes)
+read_list_file(const string& filename,
+               const map<string, string>& variables,
+               vector<FileRecord>* files,
+               vector<string>* excludes)
 {
     int err = 0;
     FILE* f = NULL;
@@ -192,11 +244,27 @@ read_list_file(const string& filename, vector<FileRecord>* files,
             
             if (words.size() == 1) {
                 // pattern: DEST
-                add_file(files, filename, i+1, words[0], words[0]);
+                bool error = false;
+                string w0 = replace_variables(words[0], variables, &error);
+                if (error) {
+                    err = 1;
+                    goto cleanup;
+                }
+                add_file(files, filename, i+1, w0, w0);
             }
             else if (words.size() == 2) {
                 // pattern: SRC DEST
-                add_file(files, filename, i+1, words[0], words[1]);
+                bool error = false;
+                string w0, w1;
+                w0 = replace_variables(words[0], variables, &error);
+                if (!error) {
+                    w1 = replace_variables(words[1], variables, &error);
+                }
+                if (error) {
+                    err = 1;
+                    goto cleanup;
+                }
+                add_file(files, filename, i+1, w0, w1);
             }
             else {
                 fprintf(stderr, "%s:%d: bad format: %s\n", filename.c_str(),
diff --git a/tools/atree/files.h b/tools/atree/files.h
index b8f0431..6480c98 100644
--- a/tools/atree/files.h
+++ b/tools/atree/files.h
@@ -1,6 +1,7 @@
 #ifndef FILES_H
 #define FILES_H
 
+#include <map>
 #include <string>
 #include <vector>
 #include <sys/types.h>
@@ -25,8 +26,10 @@ struct FileRecord
     unsigned int mode;
 };
 
-int read_list_file(const string& filename, vector<FileRecord>* files,
-                    vector<string>* excludes);
+int read_list_file(const string& filename,
+                   const map<string, string>& variables,
+                   vector<FileRecord>* files,
+                   vector<string>* excludes);
 int locate(FileRecord* rec, const vector<string>& search);
 void stat_out(const string& base, FileRecord* rec);
 string dir_part(const string& filename);