Choreo.h

// Copyright (c) Choreo contributors
 
#pragma once
 
#include <concepts>
#include <optional>
#include <string>
#include <string_view>
#include <unordered_map>
 
#include <fmt/format.h>
#include <frc/Errors.h>
#include <frc/Filesystem.h>
#include <frc2/command/Subsystem.h>
#include <hal/FRCUsageReporting.h>
#include <wpi/MemoryBuffer.h>
#include <wpi/json.h>
 
#include "choreo/trajectory/DifferentialSample.h"
#include "choreo/trajectory/SwerveSample.h"
#include "choreo/trajectory/Trajectory.h"
#include "choreo/trajectory/TrajectorySample.h"
#include "choreo/util/TrajSchemaVersion.h"
 
namespace choreo {
 
class Choreo {
 public:
  template <TrajectorySample SampleType>
  static std::optional<Trajectory<SampleType>> LoadTrajectory(
      std::string_view trajectoryName) {
    if (trajectoryName.ends_with(TRAJECTORY_FILE_EXTENSION)) {
      trajectoryName = trajectoryName.substr(
          0, trajectoryName.size() - TRAJECTORY_FILE_EXTENSION.size());
    }
 
    std::string trajectoryFileName = fmt::format(
        "{}/{}{}", CHOREO_DIR, trajectoryName, TRAJECTORY_FILE_EXTENSION);
 
    auto fileBuffer = wpi::MemoryBuffer::GetFile(trajectoryFileName);
    if (!fileBuffer) {
      FRC_ReportError(frc::warn::Warning, "Could not find trajectory file: {}",
                      trajectoryName);
      return {};
    }
 
    try {
      return LoadTrajectoryString<SampleType>(
          std::string{fileBuffer.value()->GetCharBuffer().data(),
                      fileBuffer.value()->size()},
          trajectoryName);
    } catch (wpi::json::parse_error& ex) {
      FRC_ReportError(frc::warn::Warning, "Could not parse trajectory file: {}",
                      trajectoryName);
      FRC_ReportError(frc::warn::Warning, "{}", ex.what());
      return {};
    }
    return {};
  }
 
  template <TrajectorySample SampleType>
  static std::optional<Trajectory<SampleType>> LoadTrajectoryString(
      std::string_view trajectoryJsonString, std::string_view trajectoryName) {
    if constexpr (std::same_as<SampleType, SwerveSample>) {
      HAL_Report(HALUsageReporting::kResourceType_ChoreoTrajectory, 1);
    } else if constexpr (std::same_as<SampleType, DifferentialSample>) {
      HAL_Report(HALUsageReporting::kResourceType_ChoreoTrajectory, 2);
    }
 
    wpi::json json = wpi::json::parse(trajectoryJsonString);
    uint32_t version = json["version"];
    if (version != kTrajSchemaVersion) {
      throw fmt::format("{}.traj: Wrong version {}. Expected {}",
                        trajectoryName, version, kTrajSchemaVersion);
    }
    Trajectory<SampleType> trajectory;
    from_json(json, trajectory);
    return trajectory;
  }
 
  template <TrajectorySample SampleType>
  class TrajectoryCache {
   public:
    static std::optional<Trajectory<SampleType>> LoadTrajectory(
        std::string_view trajectoryName) {
      if (cache.contains(std::string{trajectoryName})) {
        return cache[std::string{trajectoryName}];
      } else {
        cache[std::string{trajectoryName}] =
            Choreo::LoadTrajectory<SampleType>(trajectoryName);
        return cache[std::string{trajectoryName}];
      }
    }
 
    static std::optional<Trajectory<SampleType>> LoadTrajectory(
        std::string_view trajectoryName, int splitIndex) {
      std::string key = fmt::format("{}.:.{}", trajectoryName, splitIndex);
 
      if (!cache.contains(key)) {
        if (cache.contains(std::string{trajectoryName})) {
          cache[key] =
              cache[std::string{trajectoryName}].value().GetSplit(splitIndex);
        } else {
          auto possibleTrajectory = LoadTrajectory(trajectoryName);
          cache[std::string{trajectoryName}] = possibleTrajectory;
 
          if (possibleTrajectory.has_value()) {
            cache[key] = possibleTrajectory.value().GetSplit(splitIndex);
          }
        }
      }
 
      return cache[key];
    }
 
    static void Clear() { cache.clear(); }
 
   private:
    static inline std::unordered_map<std::string,
                                     std::optional<Trajectory<SampleType>>>
        cache;
  };
 
 private:
  static constexpr std::string_view TRAJECTORY_FILE_EXTENSION = ".traj";
 
  static inline const std::string CHOREO_DIR =
      frc::filesystem::GetDeployDirectory() + "/choreo";
 
  Choreo();
};
 
}  // namespace choreo