Fixed protected functions computation

Makefile

@@ -15,7 +15,7 @@ doc:
 	doxygen Doxyfile
 
 clean:
-	rm -rf build
+	rm -rf build extracted
 	rm -rf cmake-build-debug
 	rm -rf src/apex.log
 	rm -rf src/apex.out

apex.py

@@ -7,10 +7,15 @@
 
 """
 python3 apex.py examples/mod1/example_mod1.bc example_mod1.c 11 --export=true
+
+output:
+
+    extracted: binary program extracted from the example_mod1.bc
+    build/*: directory full of interesting files, logs, etc.
 """
 
 import argparse
-import shlex
+import os
 import subprocess
 import sys
 
@@ -37,7 +42,7 @@ def execute(cmd, print_output=False):
 line = args.line
 export = args.export
 
-execute("rm -rf build; mkdir build")
+execute("rm -rf build extracted; mkdir build")
 
 # Compile apexlib to the bytecode.
 execute("clang -O0 -g -c -emit-llvm src/apex/apexlib.c -o build/apexlib.bc")
@@ -50,19 +55,20 @@ def execute(cmd, print_output=False):
 
 # Run APEXPass on the linked bytecode we produced above.
 # Save log to the build/apex.log.
+if not os.path.isfile("src/build/apex/libAPEXPass.so"):
+    print("ERROR: Please first build APEX with: make build")
+    sys.exit(1)
 opt = """opt -o build/apex.bc -load src/build/apex/libAPEXPass.so -apex -file={FILE} -line={LINE} < build/linked.bc 2> build/apex.log
       """.format(FILE=target_file, LINE=line)
 execute(opt)
 
+# Compile apex.bc into executable called "extracted"
+execute("clang -o extracted build/apex.bc")
+
 # Disassembly apexlib and final extracted bytecode for dbg & logging purposes.
 execute("llvm-dis build/apexlib.bc -o build/apexlib.ll")
 execute("llvm-dis build/apex.bc -o build/apex.ll")
 
-# Run final extracted bytecode and save the result to the build/apex.out.
-execute("lli build/apex.bc > build/apex.out")
-execute("cat build/apex.out", print_output=True)
-
-
 # Optional call graphs export
 if export and export == "true":
     execute("rm -rf build/callgraphs; mkdir build/callgraphs")

examples/experiments/yes/yes.c

@@ -0,0 +1,41 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define max(a,b) \
+	({ __typeof__ (a) _a = (a); \
+	 __typeof__ (b) _b = (b); \
+	 _a > _b ? _a : _b; });
+
+int main(int argc, char* argv[]) {
+	char* output;
+	int output_size;
+	int output_len;
+
+	if (argc > 1) {
+		// Create buffer.
+		int needed_size = 2 + strlen(argv[1]);
+		for (int i = 2; i < argc; i++) {
+			needed_size += 1 + strlen(argv[i]);
+		}
+
+		output = (char*) malloc(needed_size);
+
+		// Append to buffer.
+		strcat(output, argv[1]);
+		for (int i = 2; i < argc; i++) {
+			strcat(output, " ");
+			strcat(output, argv[i]);
+		}
+		strcat(output, "\n");
+		output_len = strlen(output);
+	} else {
+		output = "y\n";
+		output_len = 2;
+	}
+
+	// Flood.
+	for(;;)
+		fwrite(output, 1, output_len, stdout);
+}
+

src/apex/apex.cpp

@@ -19,6 +19,7 @@
 bool APEXPass::runOnModule(Module &M) {
   logPrintUnderline("APEXPass START.");
 
+
   if (VERBOSE_DEBUG) {
     logPrintUnderline("Initial module dump.");
     logDumpModule(M);
@@ -30,6 +31,10 @@ bool APEXPass::runOnModule(Module &M) {
   logPrintUnderline("Locating target instructions.");
   moduleFindTargetInstructionsOrDie(M, ARG_FILE, ARG_LINE);
 
+  logPrintUnderline("Collecting protected functions.");
+  collectProtectedFunctions(M);
+
+
   logPrintUnderline(
       "Initializing dg. Calculating control and data dependencies.");
   dgInit(M);
@@ -126,11 +131,10 @@ void APEXPass::logDumpModule(const Module &M) {
 // Function utilities
 // ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
-/// Returns true if function @F is protected.
-/// That is, part of the @PROTECTED_FCNS
+/// Returns true if function @F is protected,
 /// Protected functions will not be removed at the end of the APEXPass.
 bool APEXPass::functionIsProtected(const Function *F) {
-  for (std::string fcn_id : PROTECTED_FCNS) {
+  for (std::string fcn_id : protected_functions_) {
     if (F->getGlobalIdentifier() == fcn_id) {
       return true;
     }
@@ -370,18 +374,31 @@ void APEXPass::apexDgPrintDependenciesCompact() {
 
 /// Take @node and find all data dependencies that form the chain.
 /// Store these dependencies in the @dependencies vector.
-// TODO: Refactor this function. It is copy-paste code atm.
+// TODO: Refactor this function. It is copy-paste code!
 void APEXPass::apexDgFindDataDependencies(
     LLVMNode &node, std::vector<LLVMNode *> &data_dependencies,
     std::vector<LLVMNode *> &rev_data_dependencies) {
   {
+    std::vector<LLVMNode *> visited;
     std::vector<LLVMNode *> queue;
     queue.push_back(&node);
     while (false == queue.empty()) {
       // Pop @curr and investigate it first.
       LLVMNode *curr = queue.back();
       queue.pop_back();
 
+      bool already_visited = false;
+      for (LLVMNode *visited_node : visited) {
+        if (curr == visited_node) {
+          already_visited = true;
+          break;
+        }
+      }
+      if (already_visited) {
+        continue;
+      }
+      visited.push_back(curr);
+
       // Store @curr if it is not already stored.
       bool new_dependency = true;
       for (LLVMNode *node : data_dependencies) {
@@ -407,6 +424,19 @@ void APEXPass::apexDgFindDataDependencies(
       LLVMNode *curr = queue.back();
       queue.pop_back();
 
+      std::vector<LLVMNode *> visited;
+      bool already_visited = false;
+      for (LLVMNode *visited_node : visited) {
+        if (curr == visited_node) {
+          already_visited = true;
+          break;
+        }
+      }
+      if (already_visited) {
+        continue;
+      }
+      visited.push_back(curr);
+
       // Store @curr if it is not already stored.
       bool new_dependency = true;
       for (LLVMNode *node : rev_data_dependencies) {
@@ -1208,4 +1238,21 @@ void APEXPass::stripAllDebugSymbols(Module &M) {
     llvm::stripDebugInfo(F);
   }
   logPrint("- done");
+}
+
+
+/// Goes over all functions in the module and adds functions
+/// that are only declarations to the @protected_functions_.
+/// We do not want to remove declarations.
+void APEXPass::collectProtectedFunctions(Module &M) {
+  for (const auto &F: M) {
+    if (VERBOSE_DEBUG) {
+      logPrint(F.getGlobalIdentifier());
+    }
+    if (F.isDeclaration()) {
+      logPrintDbg("- is declaration");
+      protected_functions_.push_back(F.getGlobalIdentifier());
+    }
+  }
+  logPrint("- done");
 }

src/apex/apex.h

@@ -48,22 +48,6 @@ int APEX_DONE = 0;
 /// Function LLVMNodes connected via the data dependencies.
 using DependencyBlock = std::vector<LLVMNode *>;
 
-/// Protected functions IDs. These will not be removed by APEXPass.
-std::vector<std::string> PROTECTED_FCNS = {
-    // These belong to the _apex_exit function.
-    "_apex_exit",
-    "exit",
-
-    // These belong to the _apex_extract_int function.
-    "_apex_extract_int",
-    "printf",
-
-    // LLVM stuff
-    "llvm.dbg.declare",
-    "llvm.stackrestore",
-    "llvm.stacksave",
-};
-
 /// Command line arguments for opt.
 cl::opt<std::string> ARG_FILE(
     "file", cl::desc("Filename in relative path from to the launching script."),
@@ -112,21 +96,21 @@ class APEXPass : public ModulePass {
   // Data members
   // ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
-  // Source and target functions global IDs.
+  /// Source and target functions global IDs.
   std::string source_function_id_ = "main";
   std::string target_function_id_ = ""; // Will be properly initialized later.
-  // Target instructions that correspond to the user input.
+  /// Target instructions that correspond to the user input.
   std::vector<const Instruction *> target_instructions_;
 
-  // Dependence graph: https://github.com/mchalupa/dg
+  /// Dependence graph: https://github.com/mchalupa/dg
   LLVMDependenceGraph dg_;
 
-  // Holds all the necessary info about dependencies.
+  /// Holds all the necessary info about dependencies.
   APEXDependencyGraph apex_dg_;
 
-  // @path_ is the representation of computed execution path.
-  // It holds pair of function and dependency block from function through which
-  // is the execution "flowing".
+  /// @path_ is the representation of computed execution path.
+  /// It holds pair of function and dependency block from function through which
+  /// is the execution "flowing".
   std::vector<DependencyBlock> path_;
 
   std::map<const Function *, std::vector<DependencyBlock>>
@@ -135,7 +119,18 @@ class APEXPass : public ModulePass {
   std::map<DependencyBlock, std::vector<const Function *>>
       blocks_functions_callgraph_;
 
-  // Member Functions
+    /// Protected functions IDs. These will not be removed by APEXPass.
+    std::vector<std::string> protected_functions_ = {
+            // apexlib functions
+            "_apex_exit",
+            "_apex_extract_int",
+
+            // LLVM stuff
+            "llvm.stackrestore",
+            "llvm.stacksave",
+    };
+
+  // Methods
   // ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
   // Logging utilities.
@@ -177,6 +172,7 @@ class APEXPass : public ModulePass {
   void moduleInjectExitExtract(Module &M);
   void removeUnneededStuff(Module &M);
   void stripAllDebugSymbols(Module &M);
+    void collectProtectedFunctions(Module &M);
 };
 
 /// Registering our own pass, so it can be ran via opt.