From 9cefa380796c22283d218b45d2d8230c8544ec5c Mon Sep 17 00:00:00 2001
From: Gabriel Marques <gabriel.marques@schrodinger.com>
Date: Sun, 2 Jun 2024 21:23:13 -0400
Subject: [PATCH] Print CGO table (backported from Incentive)

---
 layer0/PrintUtils.cpp | 67 +++++++++++++++++++++++++++++++++++++++++++
 layer0/PrintUtils.h   | 31 ++++++++++++++++++++
 layer1/CGO.cpp        | 26 +++++++++++++++++
 3 files changed, 124 insertions(+)

diff --git a/layer0/PrintUtils.cpp b/layer0/PrintUtils.cpp
index e69de29bb..37aa269b4 100644
--- a/layer0/PrintUtils.cpp
+++ b/layer0/PrintUtils.cpp
@@ -0,0 +1,67 @@
+#include <iostream>
+#include <iomanip>
+#include <algorithm>
+#include "PrintUtils.h"
+
+display_table_t & display_table_t::begin_row() {
+  _table.push_back(std::vector<std::string>());
+  _current_row++;
+  return *this;
+}
+
+void display_table_t::display() {
+  // Dimensions of tables
+  const size_t num_rows = _table.size();
+  const size_t num_cols = ([&] () {
+    size_t largest = 0;
+    for (auto & v : _table) {
+      largest = std::max(largest, v.size());
+    }
+    return largest;
+  })();
+
+  // Pad
+  for (auto & row : _table) {
+    while (row.size() < num_cols) {
+      row.emplace_back(" ");
+    }
+  }
+
+  // Get the largest string for each col
+  std::vector<size_t> col_sizes(num_cols);
+  for (size_t j = 0; j < num_cols; ++j) {
+    size_t largest = 0;
+    for (size_t i = 0; i < num_rows; ++i) {
+      largest = std::max(largest, _table[i][j].size());
+    }
+    col_sizes[j] = largest;
+  }
+
+
+  // construct the output
+  std::stringstream ss;
+  ss << std::left;
+  auto insert_hr = [&]() {
+    for (size_t j = 0; j < num_cols; ++j) {
+      ss << std::setw(col_sizes[j] + 3)
+      << std::setfill('-') << "+";
+    }
+    ss << "+" << std::setfill(' ') << std::endl;
+  };
+  {
+    int i = 0, j = 0;
+    for (auto & row : _table) {
+      insert_hr();
+      ss << "| ";
+      for (auto & col : row) {
+        ss << std::setw(col_sizes[j])
+           << col << " | ";
+        j++;
+      }
+      ss << std::endl;
+      i++; j = 0;
+    }
+    insert_hr();
+  }
+  std::cout << ss.str();
+}
diff --git a/layer0/PrintUtils.h b/layer0/PrintUtils.h
index e69de29bb..c6df53d4b 100644
--- a/layer0/PrintUtils.h
+++ b/layer0/PrintUtils.h
@@ -0,0 +1,31 @@
+#pragma once
+#include <type_traits>
+#include <vector>
+#include <string>
+#include <sstream>
+
+/***********************************************************************
+ * Table pretty printing utility class
+ * ---------------------------------------------------------------------
+ * Will autofill to the row with the largest column
+ ***********************************************************************/
+class display_table_t {
+public:
+  // both begins and ends a row
+  display_table_t & begin_row();
+
+  // inserts a new cell with some printable type in it
+  template <typename T>
+  display_table_t & insert_cell(T s) {
+    std::stringstream ss;
+    ss << s;
+    _table[_current_row].emplace_back(ss.str());
+    return *this;
+  }
+
+  // finalizes the table and prints it to std out
+  void display();
+private:
+  int _current_row { -1 };
+  std::vector<std::vector<std::string>> _table;
+};
diff --git a/layer1/CGO.cpp b/layer1/CGO.cpp
index c94068365..81c9c449f 100644
--- a/layer1/CGO.cpp
+++ b/layer1/CGO.cpp
@@ -7981,6 +7981,32 @@ void CGO::add_to_cgo(int op, const float * pc) {
 }
 
 void CGO::print_table() const {
+  display_table_t table;
+  table.begin_row()
+    .insert_cell("#")
+    .insert_cell("OP_CODE")
+    .insert_cell("OP_SZ")
+    .insert_cell("DATA");
+
+  unsigned j = 0;
+  for (auto it = begin(); !it.is_stop(); ++it) {
+    const auto pc = it.data();
+    const auto op_code = it.op_code();
+  
+    table.begin_row()
+      .insert_cell(++j)
+      .insert_cell(op_code)
+      .insert_cell(CGO_sz[op_code]);
+    int sz = CGO_sz[op_code];
+    std::stringstream ss;
+    for (int i = 0; i < sz; i++) {
+      ss << CGO_get_int(pc + i);
+      if (i != (sz - 1))
+        ss << ", ";
+    }
+    table.insert_cell(ss.str());
+  }
+  table.display();
 }
 
 CGO* CGOConvertSpheresToPoints(const CGO* I)