Add --yaml-stream

sparkprime · sparkprime · commit be5a809eee8f · 2016-03-01T19:03:39.000-05:00
diff --git a/cmd/jsonnet.cpp b/cmd/jsonnet.cpp
@@ -87,6 +87,7 @@ void usage(std::ostream &o)
     o << "  --code-file <var>=<val> As --file but file contents is Jsonnet code\n";
     o << "  -o / --output-file <file> Write to the output file rather than stdout\n";
     o << "  -m / --multi <dir>      Write multiple files to the directory, list files on stdout\n";
+    o << "  -y / --yaml-stream      Write output as a YAML stream of JSON documents\n";
     o << "  -S / --string           Expect a string, manifest as plain text\n";
     o << "  -s / --max-stack <n>    Number of allowed stack frames\n";
     o << "  -t / --max-trace <n>    Max length of stack trace before cropping\n";
@@ -148,6 +149,7 @@ struct JsonnetConfig {
 
     // EVAL flags
     bool evalMulti;
+    bool evalStream;
     std::string evalMultiOutputDir;
 
     // FMT flags
@@ -157,6 +159,7 @@ struct JsonnetConfig {
     JsonnetConfig()
       : cmd(EVAL), filenameIsCode(false),
         evalMulti(false),
+        evalStream(false),
         fmtInPlace(false),
         fmtTest(false)
     { }
@@ -328,6 +331,8 @@ static bool process_args(int argc,
                     output_dir += '/';
                 }
                 config->evalMultiOutputDir = output_dir;
+            } else if (arg == "-y" || arg == "--yaml-stream") {
+                config->evalStream = true;
             } else if (arg == "-S" || arg == "--string") {
                 jsonnet_string_output(vm, 1);
             } else {
@@ -457,7 +462,8 @@ static bool read_input(JsonnetConfig* config, std::string* input) {
 
 /** Writes output files for multiple file output */
 static bool write_multi_output_files(JsonnetVm* vm, char* output,
-                                     const std::string& output_dir) {
+                                     const std::string& output_dir)
+{
     // If multiple file output is used, then iterate over each string from
     // the sequence of strings returned by jsonnet_evaluate_snippet_multi,
     // construct pairs of filename and content, and write each output file.
@@ -513,6 +519,30 @@ static bool write_multi_output_files(JsonnetVm* vm, char* output,
     return true;
 }
 
+/** Writes output files for multiple file output */
+static bool write_output_stream(JsonnetVm* vm, char* output)
+{
+    // If multiple file output is used, then iterate over each string from
+    // the sequence of strings returned by jsonnet_evaluate_snippet_multi,
+    // construct pairs of filename and content, and write each output file.
+    std::vector<std::string> r;
+    for (const char *c=output ; *c!='\0' ; ) {
+        const char *json = c;
+        while (*c != '\0') ++c;
+        ++c;
+        r.emplace_back(json);
+    }
+    jsonnet_realloc(vm, output, 0);
+    for (const auto &str : r) {
+        std::cout << "---\n";
+        std::cout << str;
+    }
+    if (r.size() > 0)
+        std::cout << "...\n";
+    std::cout.flush();
+    return true;
+}
+
 /** Writes the output JSON to the specified output file for single-file
  * output
  */
@@ -563,6 +593,9 @@ int main(int argc, const char **argv)
                 if (config.evalMulti) {
                     output = jsonnet_evaluate_snippet_multi(
                         vm, config.inputFile.c_str(), input.c_str(), &error);
+                } else if (config.evalStream) {
+                    output = jsonnet_evaluate_snippet_stream(
+                        vm, config.inputFile.c_str(), input.c_str(), &error);
                 } else {
                     output = jsonnet_evaluate_snippet(
                         vm, config.inputFile.c_str(), input.c_str(), &error);
@@ -581,8 +614,12 @@ int main(int argc, const char **argv)
                     if (!write_multi_output_files(vm, output, config.evalMultiOutputDir)) {
                         return EXIT_FAILURE;
                     }
+                } else if (config.evalStream) {
+                    if (!write_output_stream(vm, output)) {
+                        return EXIT_FAILURE;
+                    }
                 } else {
-                    bool successful =  write_output_file(output, config.outputFile);
+                    bool successful = write_output_file(output, config.outputFile);
                     jsonnet_realloc(vm, output, 0);
                     if (!successful) {
                         jsonnet_destroy(vm);
diff --git a/core/libjsonnet.cpp b/core/libjsonnet.cpp
@@ -341,59 +341,90 @@ char *jsonnet_fmt_snippet(JsonnetVm *vm, const char *filename, const char *snipp
 }
 
 
+namespace {
+  enum EvalKind { REGULAR, MULTI, STREAM };
+}
+
 static char *jsonnet_evaluate_snippet_aux(JsonnetVm *vm, const char *filename,
-                                          const char *snippet, int *error, bool multi)
+                                          const char *snippet, int *error, EvalKind kind)
 {
     try {
         Allocator alloc;
-        std::string json_str;
         AST *expr;
-        std::map<std::string, std::string> files;
         Tokens tokens = jsonnet_lex(filename, snippet);
 
         expr = jsonnet_parse(&alloc, tokens);
 
         jsonnet_desugar(&alloc, expr);
 
         jsonnet_static_analysis(expr);
-        if (multi) {
-            files = jsonnet_vm_execute_multi(&alloc, expr, vm->ext, vm->maxStack,
-                                             vm->gcMinObjects, vm->gcGrowthTrigger,
-                                             vm->importCallback, vm->importCallbackContext,
-                                             vm->stringOutput);
-        } else {
-            json_str = jsonnet_vm_execute(&alloc, expr, vm->ext, vm->maxStack,
-                                          vm->gcMinObjects, vm->gcGrowthTrigger,
-                                          vm->importCallback, vm->importCallbackContext,
-                                          vm->stringOutput);
-            json_str += "\n";
-        }
-
-        if (multi) {
-            size_t sz = 1; // final sentinel
-            for (const auto &pair : files) {
-                sz += pair.first.length() + 1; // include sentinel
-                sz += pair.second.length() + 2; // Add a '\n' as well as sentinel
+        switch (kind) {
+            case REGULAR: {
+                std::string json_str = jsonnet_vm_execute(
+                    &alloc, expr, vm->ext, vm->maxStack, vm->gcMinObjects, vm->gcGrowthTrigger,
+                    vm->importCallback, vm->importCallbackContext, vm->stringOutput);
+                json_str += "\n";
+                *error = false;
+                return from_string(vm, json_str);
             }
-            char *buf = (char*)::malloc(sz);
-            if (buf == nullptr) memory_panic();
-            std::ptrdiff_t i = 0;
-            for (const auto &pair : files) {
-                memcpy(&buf[i], pair.first.c_str(), pair.first.length() + 1);
-                i += pair.first.length() + 1;
-                memcpy(&buf[i], pair.second.c_str(), pair.second.length());
-                i += pair.second.length();
-                buf[i] = '\n';
-                i++;
-                buf[i] = '\0';
-                i++;
+            break;
+
+            case MULTI: {
+                std::map<std::string, std::string> files = jsonnet_vm_execute_multi(
+                    &alloc, expr, vm->ext, vm->maxStack, vm->gcMinObjects, vm->gcGrowthTrigger,
+                    vm->importCallback, vm->importCallbackContext, vm->stringOutput);
+                size_t sz = 1; // final sentinel
+                for (const auto &pair : files) {
+                    sz += pair.first.length() + 1; // include sentinel
+                    sz += pair.second.length() + 2; // Add a '\n' as well as sentinel
+                }
+                char *buf = (char*)::malloc(sz);
+                if (buf == nullptr) memory_panic();
+                std::ptrdiff_t i = 0;
+                for (const auto &pair : files) {
+                    memcpy(&buf[i], pair.first.c_str(), pair.first.length() + 1);
+                    i += pair.first.length() + 1;
+                    memcpy(&buf[i], pair.second.c_str(), pair.second.length());
+                    i += pair.second.length();
+                    buf[i] = '\n';
+                    i++;
+                    buf[i] = '\0';
+                    i++;
+                }
+                buf[i] = '\0'; // final sentinel
+                *error = false;
+                return buf;
             }
-            buf[i] = '\0'; // final sentinel
-            *error = false;
-            return buf;
-        } else {
-            *error = false;
-            return from_string(vm, json_str);
+            break;
+
+            case STREAM: {
+                std::vector<std::string> documents = jsonnet_vm_execute_stream(
+                    &alloc, expr, vm->ext, vm->maxStack, vm->gcMinObjects, vm->gcGrowthTrigger,
+                    vm->importCallback, vm->importCallbackContext);
+                size_t sz = 1; // final sentinel
+                for (const auto &doc : documents) {
+                    sz += doc.length() + 2; // Add a '\n' as well as sentinel
+                }
+                char *buf = (char*)::malloc(sz);
+                if (buf == nullptr) memory_panic();
+                std::ptrdiff_t i = 0;
+                for (const auto &doc : documents) {
+                    memcpy(&buf[i], doc.c_str(), doc.length());
+                    i += doc.length();
+                    buf[i] = '\n';
+                    i++;
+                    buf[i] = '\0';
+                    i++;
+                }
+                buf[i] = '\0'; // final sentinel
+                *error = false;
+                return buf;
+            }
+            break;
+
+            default:
+            fputs("INTERNAL ERROR: bad value of 'kind', probably memory corruption.\n", stderr);
+            abort();
         }
 
     } catch (StaticError &e) {
@@ -423,7 +454,7 @@ static char *jsonnet_evaluate_snippet_aux(JsonnetVm *vm, const char *filename,
 
 }
 
-static char *jsonnet_evaluate_file_aux(JsonnetVm *vm, const char *filename, int *error, bool multi)
+static char *jsonnet_evaluate_file_aux(JsonnetVm *vm, const char *filename, int *error, EvalKind kind)
 {
     std::ifstream f;
     f.open(filename);
@@ -437,29 +468,37 @@ static char *jsonnet_evaluate_file_aux(JsonnetVm *vm, const char *filename, int
     input.assign(std::istreambuf_iterator<char>(f),
                  std::istreambuf_iterator<char>());
 
-    return jsonnet_evaluate_snippet_aux(vm, filename, input.c_str(), error, multi);
+    return jsonnet_evaluate_snippet_aux(vm, filename, input.c_str(), error, kind);
 }
 
 char *jsonnet_evaluate_file(JsonnetVm *vm, const char *filename, int *error)
 {
     TRY
-    return jsonnet_evaluate_file_aux(vm, filename, error, false);
+    return jsonnet_evaluate_file_aux(vm, filename, error, REGULAR);
     CATCH("jsonnet_evaluate_file")
     return nullptr;  // Never happens.
 }
 
 char *jsonnet_evaluate_file_multi(JsonnetVm *vm, const char *filename, int *error)
 {
     TRY
-    return jsonnet_evaluate_file_aux(vm, filename, error, true);
+    return jsonnet_evaluate_file_aux(vm, filename, error, MULTI);
     CATCH("jsonnet_evaluate_file_multi")
     return nullptr;  // Never happens.
 }
 
+char *jsonnet_evaluate_file_stream(JsonnetVm *vm, const char *filename, int *error)
+{
+    TRY
+    return jsonnet_evaluate_file_aux(vm, filename, error, STREAM);
+    CATCH("jsonnet_evaluate_file_stream")
+    return nullptr;  // Never happens.
+}
+
 char *jsonnet_evaluate_snippet(JsonnetVm *vm, const char *filename, const char *snippet, int *error)
 {
     TRY
-    return jsonnet_evaluate_snippet_aux(vm, filename, snippet, error, false);
+    return jsonnet_evaluate_snippet_aux(vm, filename, snippet, error, REGULAR);
     CATCH("jsonnet_evaluate_snippet")
     return nullptr;  // Never happens.
 }
@@ -468,11 +507,20 @@ char *jsonnet_evaluate_snippet_multi(JsonnetVm *vm, const char *filename,
                                      const char *snippet, int *error)
 {
     TRY
-    return jsonnet_evaluate_snippet_aux(vm, filename, snippet, error, true);
+    return jsonnet_evaluate_snippet_aux(vm, filename, snippet, error, MULTI);
     CATCH("jsonnet_evaluate_snippet_multi")
     return nullptr;  // Never happens.
 }
 
+char *jsonnet_evaluate_snippet_stream(JsonnetVm *vm, const char *filename,
+                                     const char *snippet, int *error)
+{
+    TRY
+    return jsonnet_evaluate_snippet_aux(vm, filename, snippet, error, STREAM);
+    CATCH("jsonnet_evaluate_snippet_stream")
+    return nullptr;  // Never happens.
+}
+
 char *jsonnet_realloc(JsonnetVm *vm, char *str, size_t sz)
 {
     (void) vm;
diff --git a/core/vm.cpp b/core/vm.cpp
@@ -2292,6 +2292,42 @@ namespace {
             return r;
         }
 
+        std::vector<std::string> manifestStream(void)
+        {
+            std::vector<std::string> r;
+            LocationRange loc("During manifestation");
+            if (scratch.t != Value::ARRAY) {
+                std::stringstream ss;
+                ss << "Stream mode: Top-level object was a " << type_str(scratch.t) << ", "
+                   << "should be an array whose elements hold "
+                   << "the JSON for each document in the stream.";
+                throw makeError(loc, ss.str());
+            }
+            auto *arr = static_cast<HeapArray*>(scratch.v.h);
+            for (auto *thunk : arr->elements) {
+                LocationRange tloc = thunk->body == nullptr
+                                   ? loc
+                                   : thunk->body->location;
+                if (thunk->filled) {
+                    stack.newCall(loc, thunk, nullptr, 0, BindingFrame{});
+                    // Keep arr alive when scratch is overwritten
+                    stack.top().val = scratch;
+                    scratch = thunk->content;
+                } else {
+                    stack.newCall(loc, thunk,
+                                  thunk->self, thunk->offset, thunk->upValues);
+                    // Keep arr alive when scratch is overwritten
+                    stack.top().val = scratch;
+                    evaluate(thunk->body, stack.size());
+                }
+                String element = manifestJson(tloc, true, U"");
+                scratch = stack.top().val;
+                stack.pop();
+                r.push_back(encode_utf8(element));
+            }
+            return r;
+        }
+
     };
 
 }
@@ -2324,3 +2360,14 @@ StrMap jsonnet_vm_execute_multi(Allocator *alloc, const AST *ast, const ExtMap &
     return vm.manifestMulti(string_output);
 }
 
+std::vector<std::string> jsonnet_vm_execute_stream(
+  Allocator *alloc, const AST *ast, const ExtMap &ext_vars, unsigned max_stack,
+  double gc_min_objects, double gc_growth_trigger, JsonnetImportCallback *import_callback,
+  void *ctx)
+{
+    Interpreter vm(alloc, ext_vars, max_stack, gc_min_objects, gc_growth_trigger,
+                   import_callback, ctx);
+    vm.evaluate(ast, 0);
+    return vm.manifestStream();
+}
+
diff --git a/core/vm.h b/core/vm.h
@@ -73,7 +73,7 @@ std::string jsonnet_vm_execute(Allocator *alloc, const AST *ast,
                                JsonnetImportCallback *import_callback, void *import_callback_ctx,
                                bool string_output);
 
-/** Execute the program and return the value as a number of JSON files.
+/** Execute the program and return the value as a number of named JSON files.
  *
  * This assumes the given program yields an object whose keys are filenames.
  *
@@ -95,4 +95,26 @@ std::map<std::string, std::string> jsonnet_vm_execute_multi(
     JsonnetImportCallback *import_callback, void *import_callback_ctx,
     bool string_output);
 
+/** Execute the program and return the value as a stream of JSON files.
+ *
+ * This assumes the given program yields an array whose elements are individual
+ * JSON files.
+ *
+ * \param alloc The allocator used to create the ast.
+ * \param ast The program to execute.
+ * \param ext The external vars / code.
+ * \param max_stack Recursion beyond this level gives an error.
+ * \param gc_min_objects The garbage collector does not run when the heap is this small.
+ * \param gc_growth_trigger Growth since last garbage collection cycle to trigger a new cycle.
+ * \param import_callback A callback to handle imports
+ * \param import_callback_ctx Context param for the import callback.
+ * \param output_string Whether to expect a string and output it without JSON encoding
+ * \throws RuntimeError reports runtime errors in the program.
+ * \returns A mapping from filename to the JSON strings for that file.
+ */
+std::vector<std::string> jsonnet_vm_execute_stream(
+    Allocator *alloc, const AST *ast, const std::map<std::string, VmExt> &ext,
+    unsigned max_stack, double gc_min_objects, double gc_growth_trigger,
+    JsonnetImportCallback *import_callback, void *import_callback_ctx);
+
 #endif
diff --git a/include/libjsonnet.h b/include/libjsonnet.h
