From 072e0c16e1eea0dbad53a06283de8383bf9e6766 Mon Sep 17 00:00:00 2001
From: Pavol Bielik <pbielik@acm.org>
Date: Mon, 26 Jan 2015 15:39:51 +0100
Subject: [PATCH] update rename functionality

---
 bin/js_features.js       | 15 +++++++++++++--
 lib/feature_extractor.js | 35 +++++++++++++++++++++++++++--------
 lib/output.js            |  2 +-
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/bin/js_features.js b/bin/js_features.js
index 26f13da7..386b887d 100755
--- a/bin/js_features.js
+++ b/bin/js_features.js
@@ -119,6 +119,11 @@ function processFile(file) {
 		return;
 	}
 
+	if (ARGS.print_ast){
+		console.log(output);
+		return;
+	}
+
 	if (!ARGS.nice_formatting) {
 		output = removeWhitespace(output);
 	}
@@ -155,10 +160,16 @@ function processFile(file) {
 						inferred_names[result[i].v] = result[i].inf.green;
 					}
 				}
-				console.log(UglifyJS.replaceMangled(code, file, inferred_names));
+				try {
+					var renamed_js = UglifyJS.replaceMangled(code, file, inferred_names);
+					console.log(renamed_js);
+				} catch (ex){
+					sys.error("ERROR: ".red + "failed rename '" + file + "': " + ex);
+				}
+
 			},
 			function(err) {
-				console.log("ERROR: ".red + "connecting to server '" + HOST + ":" + PORT + "' " + err);
+				sys.error("ERROR: ".red + "connecting to server '" + HOST + ":" + PORT + "' " + err);
 			});
 	}
 }
diff --git a/lib/feature_extractor.js b/lib/feature_extractor.js
index 93733bb3..2e10db67 100644
--- a/lib/feature_extractor.js
+++ b/lib/feature_extractor.js
@@ -57,14 +57,25 @@ function replaceMangled(code, file, infered_names) {
 		//replace variables with inferred names
 		stream = OutputStream({
 			beautify: true, replace_mangled: function (node) {
-				return node.definition() ? infered_names[feature_outputter.string_map.getId("$" + node.definition().id + "-" + node.name)] : node.name;
+				var label = nodeToString(node);
+				if (node.definition() && feature_outputter.string_map.hasId(label) && feature_outputter.string_map.getId(label) in infered_names){
+					return infered_names[feature_outputter.string_map.getId(label)];
+				} else {
+					return node.name;
+				}
+				//return node.definition() ? infered_names[feature_outputter.string_map.getId("$" + node.definition().id + "-" + node.name)] : node.name;
 			}
 		});
 	} else {
 		//replace variables with placeholders. Using in the online demo for interactive renaming.
 		stream = OutputStream({
 			beautify: true, replace_mangled: function (node) {
-				return node.definition() ? "local$$" + feature_outputter.string_map.getId("$" + node.definition().id + "-" + node.name) : node.name;
+				if (node.definition() && feature_outputter.string_map.hasId(nodeToString(node))){
+					return "local$$" + feature_outputter.string_map.getId(nodeToString(node));
+				} else {
+					return node.name;
+				}
+				//return node.definition() ? "local$$" + feature_outputter.string_map.getId("$" + node.definition().id + "-" + node.name) : node.name;
 			}
 		});
 	}
@@ -213,13 +224,13 @@ function printAst(toplevel){
 	var output = "";
 
 	var walker = new TreeWalker(function(node){
-		output += string_template("  node{id} [label=\"{label}\"];\n", {
+		output += string_template('  node{id} [label="{label}"];\n', {
 			id: node.id,
 			label: nodeType(node)
 		});
 
 		if (walker.parent() != null) {
-			output += string_template("  node{id1} -> node{id2} [weight=1];\n", {
+			output += string_template('  node{id1} -> node{id2} [weight=1];\n', {
 				id1: walker.parent().id,
 				id2: node.id
 			});
@@ -383,7 +394,7 @@ NodePathFinder.prototype.find = function(node) {
 /* ---[ JsonOutputter ]--- */
 
 function FeatureJsonOutputter() {
-	this.string_map = new StringMap(false);
+	this.string_map = new StringMap();
 	this.first_element = true;
 	this.output = "";
 	this.depth = 0;
@@ -540,18 +551,26 @@ FeatureJsonOutputter.prototype.endScope = function(){
 
 /* -----[ StringMap ]----- */
 
-function StringMap(nice_names) {
+function StringMap() {
 	this.map = {};
 	this.current_id = 0;
-	this.nice_names = nice_names;
 	this.keys = [];
 }
 
+StringMap.prototype.hasId = function(input){
+	if (input == null){
+		throw new Error("error null");
+	}
+
+	//we add a special character in from to allow for keys such as "toString"
+	var escaped_input = "#" + input;
+	return escaped_input in this.map;
+};
+
 StringMap.prototype.getId = function(input){
 	if (input == null){
 		throw new Error("error null");
 	}
-	if (this.nice_names) return input;
 
 	//we add a special character in from to allow for keys such as "toString"
 	var escaped_input = "#" + input;
diff --git a/lib/output.js b/lib/output.js
index fc786e98..c93e865a 100644
--- a/lib/output.js
+++ b/lib/output.js
@@ -1101,7 +1101,7 @@ function OutputStream(options) {
         var def = self.definition();
         // output.print_name(def ? def.mangled_name || def.name : self.name);
 
-        if (output.replace_mangled == null || self instanceof AST_This || self.unmangleable())
+        if (output.replace_mangled == null || !def || self instanceof AST_This || self.unmangleable())
             //retain original name for variables that cannot be mangled
             output.print_name(def ? def.mangled_name || def.name : self.name);
         else {