Add flatbuffer Matrix table and library

This makes it easier to pack Eigen matrices into flatbuffers for use in
constants files, AOS messages, etc.

Change-Id: Icd1f5d9d3e57821c2aa21eef03d52e67f80faa69
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/math/BUILD b/frc971/math/BUILD
new file mode 100644
index 0000000..54df5bd
--- /dev/null
+++ b/frc971/math/BUILD
@@ -0,0 +1,29 @@
+load("//aos/flatbuffers:generate.bzl", "static_flatbuffer")
+
+static_flatbuffer(
+    name = "matrix_fbs",
+    srcs = ["matrix.fbs"],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "flatbuffers_matrix",
+    hdrs = ["flatbuffers_matrix.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":matrix_fbs",
+        "//aos:json_to_flatbuffer",
+        "@com_github_google_glog//:glog",
+        "@com_github_tartanllama_expected",
+        "@org_tuxfamily_eigen//:eigen",
+    ],
+)
+
+cc_test(
+    name = "flatbuffers_matrix_test",
+    srcs = ["flatbuffers_matrix_test.cc"],
+    deps = [
+        ":flatbuffers_matrix",
+        "//aos/testing:googletest",
+    ],
+)
diff --git a/frc971/math/flatbuffers_matrix.h b/frc971/math/flatbuffers_matrix.h
new file mode 100644
index 0000000..57013c9
--- /dev/null
+++ b/frc971/math/flatbuffers_matrix.h
@@ -0,0 +1,133 @@
+#ifndef FRC971_MATH_FLATBUFFERS_MATRIX_H_
+#define FRC971_MATH_FLATBUFFERS_MATRIX_H_
+// This library provides utilities for converting between a frc971.fbs.Matrix
+// flatbuffer type and an Eigen::Matrix.
+// The interesting methods are ToEigen(), ToEigenOrDie(), and FromEigen().
+#include "glog/logging.h"
+#include "tl/expected.hpp"
+#include <Eigen/Core>
+
+#include "frc971/math/matrix_static.h"
+
+namespace frc971 {
+inline constexpr Eigen::StorageOptions ToEigenStorageOrder(
+    fbs::StorageOrder storage_order, int Rows, int Cols) {
+  // Eigen only implements one *Major version of the Matrix class for vectors.
+  // See https://eigen.tuxfamily.org/bz/show_bug.cgi?id=416
+  if (Rows == 1) {
+    return Eigen::RowMajor;
+  }
+  if (Cols == 1) {
+    return Eigen::ColMajor;
+  }
+  return storage_order == fbs::StorageOrder::ColMajor ? Eigen::ColMajor
+                                                      : Eigen::RowMajor;
+}
+
+template <int Rows, int Cols, fbs::StorageOrder StorageOrder>
+struct EigenMatrix {
+  typedef Eigen::Matrix<double, Rows, Cols,
+                        ToEigenStorageOrder(StorageOrder, Rows, Cols)>
+      type;
+};
+
+inline std::ostream &operator<<(std::ostream &os, fbs::MatrixField field) {
+  os << fbs::EnumNameMatrixField(field);
+  return os;
+}
+
+inline std::ostream &operator<<(std::ostream &os, fbs::FieldError error) {
+  os << fbs::EnumNameFieldError(error);
+  return os;
+}
+
+struct ConversionFailure {
+  fbs::MatrixField field;
+  fbs::FieldError error;
+  bool operator==(const ConversionFailure &) const = default;
+};
+
+inline std::ostream &operator<<(std::ostream &os, ConversionFailure failure) {
+  os << "(" << failure.field << ", " << failure.error << ")";
+  return os;
+}
+
+template <int Rows, int Cols,
+          fbs::StorageOrder StorageOrder = fbs::StorageOrder::ColMajor>
+tl::expected<typename EigenMatrix<Rows, Cols, StorageOrder>::type,
+             ConversionFailure>
+ToEigen(const fbs::Matrix &matrix) {
+  if (!matrix.has_rows()) {
+    return tl::unexpected(
+        ConversionFailure{fbs::MatrixField::kRows, fbs::FieldError::kMissing});
+  }
+  if (!matrix.has_cols()) {
+    return tl::unexpected(
+        ConversionFailure{fbs::MatrixField::kCols, fbs::FieldError::kMissing});
+  }
+  if (!matrix.has_data()) {
+    return tl::unexpected(
+        ConversionFailure{fbs::MatrixField::kData, fbs::FieldError::kMissing});
+  }
+  if (matrix.rows() != Rows) {
+    return tl::unexpected(ConversionFailure{
+        fbs::MatrixField::kRows, fbs::FieldError::kInconsistentWithTemplate});
+  }
+  if (matrix.cols() != Cols) {
+    return tl::unexpected(ConversionFailure{
+        fbs::MatrixField::kCols, fbs::FieldError::kInconsistentWithTemplate});
+  }
+  if (matrix.storage_order() != StorageOrder) {
+    return tl::unexpected(
+        ConversionFailure{fbs::MatrixField::kStorageOrder,
+                          fbs::FieldError::kInconsistentWithTemplate});
+  }
+  if (matrix.data()->size() != Rows * Cols) {
+    return tl::unexpected(ConversionFailure{
+        fbs::MatrixField::kData, fbs::FieldError::kInconsistentWithTemplate});
+  }
+  return typename EigenMatrix<Rows, Cols, StorageOrder>::type(
+      matrix.data()->data());
+}
+
+template <int Rows, int Cols,
+          fbs::StorageOrder StorageOrder = fbs::StorageOrder::ColMajor>
+typename EigenMatrix<Rows, Cols, StorageOrder>::type ToEigenOrDie(
+    const fbs::Matrix &matrix) {
+  auto result = ToEigen<Rows, Cols, StorageOrder>(matrix);
+  if (!result.has_value()) {
+    LOG(FATAL) << "Failed to convert to matrix with error " << result.error()
+               << ".";
+  }
+  return result.value();
+}
+
+template <int Rows, int Cols,
+          fbs::StorageOrder StorageOrder = fbs::StorageOrder::ColMajor>
+bool FromEigen(
+    const typename EigenMatrix<Rows, Cols, StorageOrder>::type &matrix,
+    fbs::MatrixStatic *flatbuffer) {
+  constexpr size_t kSize = Rows * Cols;
+  auto data = flatbuffer->add_data();
+  if (!data->reserve(kSize)) {
+    return false;
+  }
+  // TODO(james): Use From*() methods once they get upstreamed...
+  data->resize(kSize);
+  std::copy(matrix.data(), matrix.data() + kSize, data->data());
+  flatbuffer->set_rows(Rows);
+  flatbuffer->set_cols(Cols);
+  flatbuffer->set_storage_order(StorageOrder);
+  return true;
+}
+
+template <typename T>
+bool FromEigen(const T &matrix, fbs::MatrixStatic *flatbuffer) {
+  return FromEigen<T::RowsAtCompileTime, T::ColsAtCompileTime,
+                   (T::IsRowMajor ? fbs::StorageOrder::RowMajor
+                                  : fbs::StorageOrder::ColMajor)>(matrix,
+                                                                  flatbuffer);
+}
+
+}  // namespace frc971
+#endif  // FRC971_MATH_FLATBUFFERS_MATRIX_H_
diff --git a/frc971/math/flatbuffers_matrix_test.cc b/frc971/math/flatbuffers_matrix_test.cc
new file mode 100644
index 0000000..2807309
--- /dev/null
+++ b/frc971/math/flatbuffers_matrix_test.cc
@@ -0,0 +1,90 @@
+#include "frc971/math/flatbuffers_matrix.h"
+
+#include "gtest/gtest.h"
+
+#include "aos/json_to_flatbuffer.h"
+
+namespace frc971::testing {
+
+class FlatbuffersMatrixTest : public ::testing::Test {
+ protected:
+  template <int Rows, int Cols,
+            fbs::StorageOrder StorageOrder = fbs::StorageOrder::ColMajor>
+  tl::expected<typename EigenMatrix<Rows, Cols, StorageOrder>::type,
+               ConversionFailure>
+  ToEigen(std::string_view json) {
+    return frc971::ToEigen<Rows, Cols, StorageOrder>(
+        aos::FlatbufferDetachedBuffer<fbs::Matrix>(
+            aos::JsonToFlatbuffer<fbs::Matrix>(json))
+            .message());
+  }
+};
+
+TEST_F(FlatbuffersMatrixTest, ReadWriteMatrix) {
+  const Eigen::Matrix<double, 3, 4> expected{
+      {0, 1, 2, 3}, {4, 5, 6, 7}, {8, 9, 10, 11}};
+  aos::fbs::Builder<fbs::MatrixStatic> builder;
+  ASSERT_TRUE(FromEigen(expected, builder.get()));
+  EXPECT_EQ(
+      "{ \"rows\": 3, \"cols\": 4, \"storage_order\": \"ColMajor\", \"data\": "
+      "[ 0.0, 4.0, 8.0, 1.0, 5.0, 9.0, 2.0, 6.0, 10.0, 3.0, 7.0, 11.0 ] }",
+      aos::FlatbufferToJson(builder.AsFlatbufferSpan()));
+
+  const Eigen::Matrix<double, 3, 4> result =
+      ToEigenOrDie<3, 4>(builder->AsFlatbuffer());
+  EXPECT_EQ(expected, result);
+}
+
+TEST_F(FlatbuffersMatrixTest, ReadWriteMatrixRowMajor) {
+  const Eigen::Matrix<double, 3, 4, Eigen::StorageOptions::RowMajor> expected{
+      {0, 1, 2, 3}, {4, 5, 6, 7}, {8, 9, 10, 11}};
+  aos::fbs::Builder<fbs::MatrixStatic> builder;
+  ASSERT_TRUE(FromEigen(expected, builder.get()));
+  EXPECT_EQ(
+      "{ \"rows\": 3, \"cols\": 4, \"storage_order\": \"RowMajor\", \"data\": "
+      "[ 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0 ] }",
+      aos::FlatbufferToJson(builder.AsFlatbufferSpan()));
+
+  const Eigen::Matrix<double, 3, 4, Eigen::StorageOptions::RowMajor> result =
+      ToEigenOrDie<3, 4, fbs::StorageOrder::RowMajor>(builder->AsFlatbuffer());
+  EXPECT_EQ(expected, result);
+}
+
+class FlatbuffersMatrixParamTest
+    : public FlatbuffersMatrixTest,
+      public ::testing::WithParamInterface<
+          std::tuple<std::string, ConversionFailure>> {};
+TEST_P(FlatbuffersMatrixParamTest, ConversionFailures) {
+  auto result = this->ToEigen<3, 4>(std::get<0>(GetParam()));
+  EXPECT_FALSE(result.has_value());
+  EXPECT_EQ(std::get<1>(GetParam()), result.error());
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ConversionFailureTests, FlatbuffersMatrixParamTest,
+    ::testing::Values(
+        std::make_tuple("{}", ConversionFailure{fbs::MatrixField::kRows,
+                                                fbs::FieldError::kMissing}),
+        std::make_tuple(R"json({"rows": 3})json",
+                        ConversionFailure{fbs::MatrixField::kCols,
+                                          fbs::FieldError::kMissing}),
+        std::make_tuple(R"json({"rows": 3, "cols": 4})json",
+                        ConversionFailure{fbs::MatrixField::kData,
+                                          fbs::FieldError::kMissing}),
+        std::make_tuple(
+            R"json({"rows": 1, "cols": 4, "data": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]})json",
+            ConversionFailure{fbs::MatrixField::kRows,
+                              fbs::FieldError::kInconsistentWithTemplate}),
+        std::make_tuple(
+            R"json({"rows": 3, "cols": 7, "data": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]})json",
+            ConversionFailure{fbs::MatrixField::kCols,
+                              fbs::FieldError::kInconsistentWithTemplate}),
+        std::make_tuple(R"json({"rows": 3, "cols": 4, "data": []})json",
+                        ConversionFailure{
+                            fbs::MatrixField::kData,
+                            fbs::FieldError::kInconsistentWithTemplate}),
+        std::make_tuple(
+            R"json({"rows": 3, "cols": 4, "storage_order": "RowMajor", "data": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]})json",
+            ConversionFailure{fbs::MatrixField::kStorageOrder,
+                              fbs::FieldError::kInconsistentWithTemplate})));
+}  // namespace frc971::testing
diff --git a/frc971/math/matrix.fbs b/frc971/math/matrix.fbs
new file mode 100644
index 0000000..8bbd03a
--- /dev/null
+++ b/frc971/math/matrix.fbs
@@ -0,0 +1,39 @@
+namespace frc971.fbs;
+
+enum StorageOrder : ubyte {
+  // Column-major; i.e., for a matrix
+  // [1 2]
+  // [3 4]
+  // The memory layout will be 1 3 2 4.
+  ColMajor = 0,
+  // Row-major; i.e., for a matrix
+  // [1 2]
+  // [3 4]
+  // The memory layout will be 1 2 3 4.
+  RowMajor = 1,
+}
+
+// Represents a dynamically-sized 2-D matrix that is either row-major or column-major.
+table Matrix {
+  // rows and cols must both be greater than zero.
+  rows:uint (id: 0);
+  cols:uint (id: 1);
+  storage_order:StorageOrder = ColMajor (id: 2);
+  // data must be present and must have a length of rows * cols.
+  data:[double] (id: 3);
+}
+
+// The below enums are used in C++ code for communicating errors in parsing
+// the matrix; they are mostly only defined in the fbs file so that we get
+// pre-generated functions for converting the enum values to strings.
+enum MatrixField : ubyte {
+  kRows = 0,
+  kCols,
+  kStorageOrder,
+  kData,
+}
+
+enum FieldError : ubyte {
+  kInconsistentWithTemplate = 0,
+  kMissing,
+}