Allow Control Loops to use Static Flatbuffers for Output and Status

Signed-off-by: Nikolai Sohmers <nikolai@sohmers.com>
Change-Id: I360add2d4b51f48d46eb3915e5696d12f81edd05
diff --git a/frc971/control_loops/control_loop.h b/frc971/control_loops/control_loop.h
index c7b8a6d..e9892a3 100644
--- a/frc971/control_loops/control_loop.h
+++ b/frc971/control_loops/control_loop.h
@@ -6,6 +6,7 @@
 
 #include "aos/actions/actor.h"
 #include "aos/events/event_loop.h"
+#include "aos/flatbuffers/static_table.h"
 #include "aos/time/time.h"
 #include "aos/util/log_interval.h"
 #include "frc971/input/joystick_state_generated.h"
@@ -17,6 +18,56 @@
 constexpr ::std::chrono::nanoseconds kLoopFrequency =
     aos::common::actions::kLoopFrequency;
 
+// In order to support using both "raw" and the static flatbuffer APIs with the
+// outputs of the ControlLoop class (i.e., the Output and Status messages), we
+// need to provide different compile-time code for each option. In order to do
+// this, we set up two different specializations of the BuilderType class, and
+// will select the appropriate one depending on the type of the flatbuffer. The
+// methods and types within BuilderType are then used within the ControlLoop
+// class whenever we need to call methods associated with the flatbuffers or
+// senders themselves.
+template <typename T, class Enable = void>
+struct BuilderType;
+
+template <class GoalType, class PositionType, class StatusType,
+          class OutputType>
+class ControlLoop;
+
+template <typename T>
+struct BuilderType<T, typename std::enable_if_t<
+                          std::is_base_of<flatbuffers::Table, T>::value>> {
+  typedef aos::Sender<T>::Builder Builder;
+  typedef T FlatbufferType;
+  static Builder MakeBuilder(aos::Sender<T> *sender) {
+    return sender->MakeBuilder();
+  }
+  template <class GoalType, class PositionType, class StatusType,
+            class OutputType>
+  static void SendZeroFlatbuffer(
+      ControlLoop<GoalType, PositionType, StatusType, OutputType> *loop,
+      Builder *builder);
+  static void CheckSent(Builder *builder) { builder->CheckSent(); }
+};
+
+template <typename T>
+struct BuilderType<
+    T, typename std::enable_if_t<std::is_base_of<aos::fbs::Table, T>::value>> {
+  typedef aos::Sender<T>::StaticBuilder Builder;
+  typedef T::Flatbuffer FlatbufferType;
+  static Builder MakeBuilder(aos::Sender<T> *sender) {
+    return sender->MakeStaticBuilder();
+  }
+  template <class GoalType, class PositionType, class StatusType,
+            class OutputType>
+  static void SendZeroFlatbuffer(
+      ControlLoop<GoalType, PositionType, StatusType, OutputType> *loop,
+      Builder *builder);
+  static void CheckSent(Builder *builder) {
+    // TODO (niko) create a CheckSent() function for static flatbuffers
+    (void)builder;
+  }
+};
+
 // Provides helper methods to assist in writing control loops.
 // It will then call the RunIteration method every cycle that it has enough
 // valid data for the control loop to run.
@@ -24,6 +75,10 @@
           class OutputType>
 class ControlLoop {
  public:
+  using StatusBuilder = typename BuilderType<StatusType>::Builder;
+  using OutputBuilder = typename BuilderType<OutputType>::Builder;
+  using OutputFlatbufferType = typename BuilderType<OutputType>::FlatbufferType;
+
   ControlLoop(aos::EventLoop *event_loop, const ::std::string &name)
       : event_loop_(event_loop), name_(name) {
     output_sender_ = event_loop_->MakeSender<OutputType>(name_);
@@ -63,9 +118,9 @@
   // Sets the output to zero.
   // Override this if a value of zero (or false) is not "off" for this
   // subsystem.
-  virtual flatbuffers::Offset<OutputType> Zero(
-      typename ::aos::Sender<OutputType>::Builder *builder) {
-    return builder->template MakeBuilder<OutputType>().Finish();
+  virtual flatbuffers::Offset<OutputFlatbufferType> Zero(
+      typename ::aos::Sender<OutputFlatbufferType>::Builder *builder) {
+    return builder->template MakeBuilder<OutputFlatbufferType>().Finish();
   }
 
  protected:
@@ -87,10 +142,8 @@
   // output is going to be ignored and set to 0.
   // status is the status of the control loop.
   // Both output and status should be filled in by the implementation.
-  virtual void RunIteration(
-      const GoalType *goal, const PositionType *position,
-      typename ::aos::Sender<OutputType>::Builder *output,
-      typename ::aos::Sender<StatusType>::Builder *status) = 0;
+  virtual void RunIteration(const GoalType *goal, const PositionType *position,
+                            OutputBuilder *output, StatusBuilder *status) = 0;
 
  private:
   static constexpr ::std::chrono::milliseconds kStaleLogInterval =
@@ -132,6 +185,28 @@
       SimpleLogInterval(kStaleLogInterval, ERROR, "no goal");
 };
 
+template <typename T>
+template <class GoalType, class PositionType, class StatusType,
+          class OutputType>
+void BuilderType<T, typename std::enable_if_t<
+                        std::is_base_of<flatbuffers::Table, T>::value>>::
+    SendZeroFlatbuffer(
+        ControlLoop<GoalType, PositionType, StatusType, OutputType> *loop,
+        Builder *builder) {
+  builder->CheckOk(builder->Send(loop->Zero(builder)));
+}
+
+template <typename T>
+template <class GoalType, class PositionType, class StatusType,
+          class OutputType>
+void BuilderType<
+    T, typename std::enable_if_t<std::is_base_of<aos::fbs::Table, T>::value>>::
+    SendZeroFlatbuffer(ControlLoop<GoalType, PositionType, StatusType,
+                                   OutputType> * /* loop */,
+                       Builder *builder) {
+  builder->CheckOk(builder->Send());
+}
+
 }  // namespace frc971::controls
 
 #include "frc971/control_loops/control_loop-tmpl.h"  // IWYU pragma: export