llvm: Deduplicate LLVM-IR strings.

This changes the format of the LLVM-IR program graphs to store a list
of unique strings, rather than LLVM-IR strings in each node. We use a
graph-level "strings" feature to store a list of the original LLVM-IR
string corresponding to each graph nodes. This allows to us to refer
to the same string from multiple nodes without duplication.

This breaks compatability with the inst2vec encoder on program graphs
generated prior to this commit.

Author: ChrisCummins

programl/graph/program_graph_builder.h

@@ -103,6 +103,9 @@ class ProgramGraphBuilder {
   inline Edge* AddEdge(const Edge::Flow& flow, int32_t position, const Node* source,
                        const Node* target);
 
+  // Return a mutable pointer to the root node in the graph.
+  Node* GetMutableRootNode() { return graph_.mutable_node(0); }
+
   // Return a mutable pointer to the graph protocol buffer.
   ProgramGraph* GetMutableProgramGraph() { return &graph_; }
 

programl/ir/llvm/inst2vec_encoder.py

@@ -43,11 +43,10 @@
 )
 
 
-def NodeFullText(node: node_pb2.Node) -> str:
+def NodeFullText(graph: program_graph_pb2.ProgramGraph, node: node_pb2.Node) -> str:
     """Get the full text of a node, or an empty string if not set."""
-    if len(node.features.feature["full_text"].bytes_list.value):
-        return node.features.feature["full_text"].bytes_list.value[0].decode("utf-8")
-    return ""
+    idx = node.features.feature["llvm_string"].int64_list.value[0]
+    return graph.features.feature["strings"].bytes_list.value[idx].decode("utf-8")
 
 
 class Inst2vecEncoder(object):
@@ -85,7 +84,7 @@ def Encode(
         """
         # Gather the instruction texts to pre-process.
         lines = [
-            [NodeFullText(node)]
+            [NodeFullText(proto, node)]
             for node in proto.node
             if node.type == node_pb2.Node.INSTRUCTION
         ]

programl/ir/llvm/inst2vec_encoder_test.py

@@ -55,10 +55,21 @@ def AddVariable(self, full_text: str):
 
     def Build(self):
         proto = super(Inst2vecGraphBuilder, self).Build()
+
+        # Add the root node string feature.
+        proto.node[0].features.feature["llvm_string"].int64_list.value[:] = [0]
+
+        # Build the strings list.
+        strings_list = list(set(self.full_texts.values()))
+        proto.features.feature["strings"].bytes_list.value[:] = [
+            string.encode("utf-8") for string in strings_list
+        ]
+
+        # Add the string indices.
         for node, full_text in self.full_texts.items():
-            proto.node[node].features.feature["full_text"].bytes_list.value.append(
-                full_text.encode("utf-8")
-            )
+            idx = strings_list.index(full_text)
+            node_feature = proto.node[node].features.feature["llvm_string"]
+            node_feature.int64_list.value.append(idx)
         return proto
 
 

programl/ir/llvm/internal/program_graph_builder.cc

@@ -39,6 +39,16 @@ namespace ir {
 namespace llvm {
 namespace internal {
 
+ProgramGraphBuilder::ProgramGraphBuilder(const ProgramGraphOptions& options)
+    : programl::graph::ProgramGraphBuilder(),
+      options_(options),
+      blockCount_(0),
+      stringsList_((*GetMutableProgramGraph()->mutable_features()->mutable_feature())["strings"]
+                       .mutable_bytes_list()) {
+  // Add an empty
+  graph::AddScalarFeature(GetMutableRootNode(), "llvm_string", AddString(""));
+}
+
 labm8::StatusOr<BasicBlockEntryExit> ProgramGraphBuilder::VisitBasicBlock(
     const ::llvm::BasicBlock& block, const Function* functionMessage, InstructionMap* instructions,
     ArgumentConsumerMap* argumentConsumers, std::vector<DataEdge>* dataEdgesToAdd) {
@@ -184,7 +194,7 @@ labm8::StatusOr<FunctionEntryExits> ProgramGraphBuilder::VisitFunction(
 
   if (function.isDeclaration()) {
     Node* node = AddInstruction("; undefined function", functionMessage);
-    graph::AddScalarFeature(node, "full_text", "");
+    graph::AddScalarFeature(node, "llvm_string", AddString(""));
     functionEntryExits.first = node;
     functionEntryExits.second.push_back(node);
     return functionEntryExits;
@@ -305,7 +315,7 @@ Node* ProgramGraphBuilder::AddLlvmInstruction(const ::llvm::Instruction* instruc
   const LlvmTextComponents text = textEncoder_.Encode(instruction);
   Node* node = AddInstruction(text.opcode_name, function);
   node->set_block(blockCount_);
-  graph::AddScalarFeature(node, "full_text", text.text);
+  graph::AddScalarFeature(node, "llvm_string", AddString(text.text));
 
   // Add profiling information features, if available.
   uint64_t profTotalWeight;
@@ -327,7 +337,7 @@ Node* ProgramGraphBuilder::AddLlvmVariable(const ::llvm::Instruction* operand,
   const LlvmTextComponents text = textEncoder_.Encode(operand);
   Node* node = AddVariable(text.lhs_type, function);
   node->set_block(blockCount_);
-  graph::AddScalarFeature(node, "full_text", text.lhs);
+  graph::AddScalarFeature(node, "llvm_string", AddString(text.lhs));
 
   return node;
 }
@@ -337,7 +347,7 @@ Node* ProgramGraphBuilder::AddLlvmVariable(const ::llvm::Argument* argument,
   const LlvmTextComponents text = textEncoder_.Encode(argument);
   Node* node = AddVariable(text.lhs_type, function);
   node->set_block(blockCount_);
-  graph::AddScalarFeature(node, "full_text", text.lhs);
+  graph::AddScalarFeature(node, "llvm_string", AddString(text.lhs));
 
   return node;
 }
@@ -346,7 +356,7 @@ Node* ProgramGraphBuilder::AddLlvmConstant(const ::llvm::Constant* constant) {
   const LlvmTextComponents text = textEncoder_.Encode(constant);
   Node* node = AddConstant(text.lhs_type);
   node->set_block(blockCount_);
-  graph::AddScalarFeature(node, "full_text", text.text);
+  graph::AddScalarFeature(node, "llvm_string", AddString(text.text));
 
   return node;
 }
@@ -436,6 +446,27 @@ void ProgramGraphBuilder::Clear() {
   programl::graph::ProgramGraphBuilder::Clear();
 }
 
+Node* ProgramGraphBuilder::GetOrCreateType(const ::llvm::Type* type) {
+  auto it = types_.find(type);
+  if (it == types_.end()) {
+    Node* node = AddLlvmType(type);
+    types_[type] = node;
+    return node;
+  }
+  return it->second;
+}
+
+int32_t ProgramGraphBuilder::AddString(const string& text) {
+  auto it = stringsListPositions_.find(text);
+  if (it == stringsListPositions_.end()) {
+    int32_t index = stringsListPositions_.size();
+    stringsListPositions_[text] = index;
+    stringsList_->add_value(text);
+    return index;
+  }
+  return it->second;
+}
+
 }  // namespace internal
 }  // namespace llvm
 }  // namespace ir

programl/ir/llvm/internal/program_graph_builder.h

@@ -64,8 +64,7 @@ using ArgumentConsumerMap =
 // A specialized program graph builder for LLVM-IR.
 class ProgramGraphBuilder : public programl::graph::ProgramGraphBuilder {
  public:
-  explicit ProgramGraphBuilder(const ProgramGraphOptions& options)
-      : programl::graph::ProgramGraphBuilder(options), blockCount_(0) {}
+  explicit ProgramGraphBuilder(const ProgramGraphOptions& options);
 
   [[nodiscard]] labm8::StatusOr<ProgramGraph> Build(const ::llvm::Module& module);
 
@@ -87,6 +86,13 @@ class ProgramGraphBuilder : public programl::graph::ProgramGraphBuilder {
   Node* AddLlvmVariable(const ::llvm::Argument* argument, const Function* function);
   Node* AddLlvmConstant(const ::llvm::Constant* constant);
 
+  // Add a string to the strings list and return its position.
+  //
+  // We use a graph-level "strings" feature to store a list of the original
+  // LLVM-IR string corresponding to each graph nodes. This allows to us to
+  // refer to the same string from multiple nodes without duplication.
+  int32_t AddString(const string& text);
+
  private:
   TextEncoder textEncoder_;
 
@@ -100,6 +106,12 @@ class ProgramGraphBuilder : public programl::graph::ProgramGraphBuilder {
   // populated by VisitBasicBlock() and consumed once all functions have been
   // visited.
   absl::flat_hash_map<const ::llvm::Constant*, std::vector<PositionalNode>> constants_;
+
+  // A mapping from string table value to its position in the "strings_table"
+  // graph-level feature.
+  absl::flat_hash_map<string, int32_t> stringsListPositions_;
+  // The underlying storage for the strings table.
+  BytesList* stringsList_;
 };
 
 }  // namespace internal

programl/ir/llvm/py/llvm_test.py

@@ -37,8 +37,10 @@
 """
 
 
-def GetStringScalar(proto, name):
-    return proto.features.feature[name].bytes_list.value[0].decode("utf-8")
+def NodeFullText(graph: program_graph_pb2.ProgramGraph, node: node_pb2.Node) -> str:
+    """Get the full text of a node, or an empty string if not set."""
+    idx = node.features.feature["llvm_string"].int64_list.value[0]
+    return graph.features.feature["strings"].bytes_list.value[idx].decode("utf-8")
 
 
 def test_simple_ir():
@@ -56,25 +58,25 @@ def test_simple_ir():
 
     assert proto.node[1].text == "add"
     assert proto.node[1].type == node_pb2.Node.INSTRUCTION
-    assert GetStringScalar(proto.node[1], "full_text") == "%3 = add nsw i32 %1, %0"
+    assert NodeFullText(proto, proto.node[1]) == "%3 = add nsw i32 %1, %0"
 
     assert proto.node[2].text == "ret"
     assert proto.node[2].type == node_pb2.Node.INSTRUCTION
-    assert GetStringScalar(proto.node[2], "full_text") == "ret i32 %3"
+    assert NodeFullText(proto, proto.node[2]) == "ret i32 %3"
 
     assert proto.node[3].text == "i32"
     assert proto.node[3].type == node_pb2.Node.VARIABLE
-    assert GetStringScalar(proto.node[3], "full_text") == "i32 %3"
+    assert NodeFullText(proto, proto.node[3]) == "i32 %3"
 
     # Use startswith() to compare names for these last two variables as thier
     # order may differ.
     assert proto.node[4].text == "i32"
     assert proto.node[4].type == node_pb2.Node.VARIABLE
-    assert GetStringScalar(proto.node[4], "full_text").startswith("i32 %")
+    assert NodeFullText(proto, proto.node[4]).startswith("i32 %")
 
     assert proto.node[5].text == "i32"
     assert proto.node[5].type == node_pb2.Node.VARIABLE
-    assert GetStringScalar(proto.node[5], "full_text").startswith("i32 %")
+    assert NodeFullText(proto, proto.node[5]).startswith("i32 %")
 
 
 def test_opt_level():