1 files changed, 121 insertions, 36 deletions

diff --git a/mrbgems/mruby-bin-strip/tools/mruby-strip/mruby-strip.c b/mrbgems/mruby-bin-strip/tools/mruby-strip/mruby-strip.c
index dee3e0cd6..0c632b7c1 100644
--- a/mrbgems/mruby-bin-strip/tools/mruby-strip/mruby-strip.c
+++ b/mrbgems/mruby-bin-strip/tools/mruby-strip/mruby-strip.c
@@ -1,58 +1,143 @@
-#include <stdio.h>
+#include <mruby.h>
+
+#ifdef MRB_NO_STDIO
+# error mruby-bin-strip conflicts 'MRB_NO_STDIO' in your build configuration
+#endif
+
 #include <stdlib.h>
-#include "mruby.h"
-#include "mruby/irep.h"
-#include "mruby/dump.h"
+#include <string.h>
+#include <mruby/irep.h>
+#include <mruby/dump.h>
 
-int
-main(int argc, char **argv)
+struct strip_args {
+  int argc_start;
+  int argc;
+  char **argv;
+  mrb_bool lvar;
+};
+
+static void
+print_usage(const char *f)
 {
-  int i, dump_result;
-  FILE **files;
-  mrb_irep **ireps;
-  mrb_state *mrb;
+  printf("Usage: %s [switches] irepfiles\n", f);
+  printf("switches:\n");
+  printf("  -l, --lvar   remove LVAR section too.\n");
+}
 
-  if (argc <= 1) {
-    fprintf(stderr, "no files to strip\n");
-    return EXIT_FAILURE;
-  }
+static int
+parse_args(int argc, char **argv, struct strip_args *args)
+{
+  int i;
 
-  files = (FILE**)malloc(sizeof(FILE*) * argc);
-  for (i = 1; i < argc; ++i) {
-    files[i] = fopen(argv[i], "rb");
+  args->argc_start = 0;
+  args->argc = argc;
+  args->argv = argv;
+  args->lvar = FALSE;
 
-    if (!files[i]) {
-      fprintf(stderr, "can't open file %s\n", argv[i]);
-      return EXIT_FAILURE;
+  for (i = 1; i < argc; ++i) {
+    const size_t len = strlen(argv[i]);
+    if (len >= 2 && argv[i][0] == '-') {
+      switch (argv[i][1]) {
+      case 'l':
+        args->lvar = TRUE;
+        break;
+      case '-':
+        if (strncmp((*argv) + 2, "lvar", len) == 0) {
+          args->lvar = TRUE;
+          break;
+        }
+      default:
+        return -1;
+      }
+    }
+    else {
+      break;
     }
   }
 
-  mrb = mrb_open();
+  args->argc_start = i;
+  return i;
+}
 
-  ireps = (mrb_irep**)malloc(sizeof(mrb_irep*) * argc);
-  for (i = 1; i < argc; ++i) {
-    ireps[i] = mrb_read_irep_file(mrb, files[i]);
-    if (!ireps[i]) {
-      fprintf(stderr, "can't read irep file %s\n", argv[i]);
+static int
+strip(mrb_state *mrb, struct strip_args *args)
+{
+  int i;
+
+  for (i = args->argc_start; i < args->argc; ++i) {
+    char *filename;
+    FILE *rfile;
+    mrb_irep *irep;
+    FILE *wfile;
+    int dump_result;
+
+    filename = args->argv[i];
+    rfile = fopen(filename, "rb");
+    if (rfile == NULL) {
+      fprintf(stderr, "can't open file for reading %s\n", filename);
       return EXIT_FAILURE;
     }
-    fclose(files[i]);
-    files[i] = fopen(argv[i], "wb");
-    if (!ireps[i]) {
-      fprintf(stderr, "can't reopen irep file %s\n", argv[i]);
+
+    irep = mrb_read_irep_file(mrb, rfile);
+    fclose(rfile);
+    if (irep == NULL) {
+      fprintf(stderr, "can't read irep file %s\n", filename);
       return EXIT_FAILURE;
     }
-  }
 
-  for (i = 1; i < argc; ++i) {
-    /* debug flag must be alway false */
-    dump_result = mrb_dump_irep_binary(mrb, ireps[i], FALSE, files[i]);
+    /* clear lv if --lvar is enabled */
+    if (args->lvar) {
+      mrb_irep_remove_lv(mrb, irep);
+    }
+
+    wfile = fopen(filename, "wb");
+    if (wfile == NULL) {
+      fprintf(stderr, "can't open file for writing %s\n", filename);
+      mrb_irep_decref(mrb, irep);
+      return EXIT_FAILURE;
+    }
+
+    /* debug flag must always be false */
+    dump_result = mrb_dump_irep_binary(mrb, irep, FALSE, wfile);
+
+    fclose(wfile);
+    mrb_irep_decref(mrb, irep);
+
     if (dump_result != MRB_DUMP_OK) {
-      fprintf(stderr, "error occur when dumping %s", argv[i]);
+      fprintf(stderr, "error occurred during dumping %s\n", filename);
       return EXIT_FAILURE;
     }
   }
+  return EXIT_SUCCESS;
+}
+
+int
+main(int argc, char **argv)
+{
+  struct strip_args args;
+  int args_result;
+  mrb_state *mrb;
+  int ret;
+
+  if (argc <= 1) {
+    printf("no files to strip\n");
+    print_usage(argv[0]);
+    return EXIT_FAILURE;
+  }
+
+  args_result = parse_args(argc, argv, &args);
+  if (args_result < 0) {
+    print_usage(argv[0]);
+    return EXIT_FAILURE;
+  }
+  mrb = mrb_open_core(mrb_default_allocf, NULL);
+  if (mrb == NULL) {
+    fputs("Invalid mrb_state, exiting mruby-strip\n", stderr);
+    return EXIT_FAILURE;
+  }
+
+  ret = strip(mrb, &args);
 
   mrb_close(mrb);
-  return EXIT_SUCCESS;
+  return ret;
 }