CodeGen: rewrite implementation [3/3]

sandbox/linux/seccomp-bpf/codegen.h

Index: sandbox/linux/seccomp-bpf/codegen.h
diff --git a/sandbox/linux/seccomp-bpf/codegen.h b/sandbox/linux/seccomp-bpf/codegen.h
index 7061654b65f9acc772eb9c41137e59cd92d12116..79f9926807939a10b9687bfeefcb964e67a9b6d6 100644
--- a/sandbox/linux/seccomp-bpf/codegen.h
+++ b/sandbox/linux/seccomp-bpf/codegen.h
@@ -5,24 +5,19 @@
 #ifndef SANDBOX_LINUX_SECCOMP_BPF_CODEGEN_H__
 #define SANDBOX_LINUX_SECCOMP_BPF_CODEGEN_H__
 
+#include <stddef.h>
 #include <stdint.h>
 
 #include <map>
 #include <vector>
 
+#include "base/macros.h"
+#include "base/tuple.h"
 #include "sandbox/sandbox_export.h"
 
 struct sock_filter;
 
 namespace sandbox {
-struct BasicBlock;
-struct Instruction;
-
-typedef std::vector<Instruction*> Instructions;
-typedef std::vector<BasicBlock*> BasicBlocks;
-typedef std::map<const Instruction*, int> BranchTargets;
-typedef std::map<const Instruction*, BasicBlock*> TargetsToBlocks;
-typedef std::map<const BasicBlock*, int> IncomingBranches;
 
 // The code generator instantiates a basic compiler that can convert a
 // graph of BPF instructions into a well-formed stream of BPF instructions.
@@ -62,89 +57,55 @@ class SANDBOX_EXPORT CodeGen {
 
   // Addr represents a logical location within the program being compiled.
   // Identical code suffixes may be collapsed to a single location.
-  using Addr = Instruction*;
+  using Addr = Program::size_type;
 
   // kNullAddr represents the "null" address value; i.e., the reserved
   // address value guaranteed to not equal any actual address values.
-  static const Addr kNullAddr;
+  static const Addr kNullAddr = -1;
 
   CodeGen();
   ~CodeGen();
 
-  // Create a new instruction. Instructions form a DAG. The instruction objects
-  // are owned by the CodeGen object. They do not need to be explicitly
-  // deleted.
+  // MakeInstruction creates a new node in the instruction DAG and
+  // returns the node's address, or returns the address of an existing
+  // equivalent node if one exists.
+  //
   // For details on the possible parameters refer to <linux/filter.h>
+  // Note: Callers need not worry about explicitly creating "jump
+  // always" instructions (i.e., BPF_JMP+BPF_JA), as CodeGen inserts
+  // these as needed.
   Addr MakeInstruction(uint16_t code,
                        uint32_t k,
                        Addr jt = kNullAddr,
                        Addr jf = kNullAddr);
 
-  // Compiles the graph of instructions into a BPF program that can be passed
-  // to the kernel. Please note that this function modifies the graph in place
-  // and must therefore only be called once per graph.
-  void Compile(Addr head, Program* program);
+  // Compile linearizes the instruction DAG into an equivalent flat
+  // BPF program that can be passed to the kernel.
+  void Compile(Addr instructions, Program* program);
 
  private:
-  friend class CodeGenUnittestHelper;
-
-  // Find all the instructions that are the target of BPF_JMPs.
-  void FindBranchTargets(const Instruction& instructions,
-                         BranchTargets* branch_targets);
-
-  // Combine instructions between "head" and "tail" into a new basic block.
-  // Basic blocks are defined as sequences of instructions whose only branch
-  // target is the very first instruction; furthermore, any BPF_JMP or BPF_RET
-  // instruction must be at the very end of the basic block.
-  BasicBlock* MakeBasicBlock(Instruction* head, Instruction* tail);
-
-  // Creates a basic block and adds it to "basic_blocks"; sets "first_block"
-  // if it is still NULL.
-  void AddBasicBlock(Instruction* head,
-                     Instruction* tail,
-                     const BranchTargets& branch_targets,
-                     TargetsToBlocks* basic_blocks,
-                     BasicBlock** first_block);
-
-  // Cuts the DAG of instructions into basic blocks.
-  BasicBlock* CutGraphIntoBasicBlocks(Instruction* instructions,
-                                      const BranchTargets& branch_targets,
-                                      TargetsToBlocks* blocks);
-
-  // Find common tail sequences of basic blocks and coalesce them.
-  void MergeTails(TargetsToBlocks* blocks);
-
-  // For each basic block, compute the number of incoming branches.
-  void ComputeIncomingBranches(BasicBlock* block,
-                               const TargetsToBlocks& targets_to_blocks,
-                               IncomingBranches* incoming_branches);
-
-  // Topologically sort the basic blocks so that all jumps are forward jumps.
-  // This is a requirement for any well-formed BPF program.
-  void TopoSortBasicBlocks(BasicBlock* first_block,
-                           const TargetsToBlocks& blocks,
-                           BasicBlocks* basic_blocks);
-
-  // Convert jt_ptr_ and jf_ptr_ fields in BPF_JMP instructions to valid
-  // jt_ and jf_ jump offsets. This can result in BPF_JA instructions being
-  // inserted, if we need to jump over more than 256 instructions.
-  void ComputeRelativeJumps(BasicBlocks* basic_blocks,
-                            const TargetsToBlocks& targets_to_blocks);
-
-  // Concatenate instructions from all basic blocks into a BPF program that
-  // can be passed to the kernel.
-  void ConcatenateBasicBlocks(const BasicBlocks&, Program* program);
-
-  // We stick all instructions and basic blocks into pools that get destroyed
-  // when the CodeGen object is destroyed. This way, we neither need to worry
-  // about explicitly managing ownership, nor do we need to worry about using
-  // smart pointers in the presence of circular references.
-  Instructions instructions_;
-  BasicBlocks basic_blocks_;
-
-  // Compile() must only ever be called once as it makes destructive changes
-  // to the DAG.
-  bool compiled_;
+  using CacheKey = Tuple4<uint16_t, uint32_t, Addr, Addr>;
+  struct CacheKeyLess {
+    bool operator()(const CacheKey& lhs, const CacheKey& rhs) const;
+  };
+
+  // Offset returns how many instructions exist in |program_| after |target|.
+  size_t Offset(Addr target) const;
+
+  // Jump appends a jump-always instruction.
+  Addr Jump(Addr next);
+
+  // Append adds a new instruction to the end of |program_|.
+  Addr Append(uint16_t code, uint32_t k, size_t jt, size_t jf);
+
+  void CheckValid(Addr addr) const;
+
+  // NOTE: program_ is the compiled program in *reverse*.
+  Program program_;
+
+  std::map<CacheKey, Addr, CacheKeyLess> cache_;
+
+  DISALLOW_COPY_AND_ASSIGN(CodeGen);
 };
 
 }  // namespace sandbox