From d5b81ec9ecfab3676b9263c02ea8b76e9d6c4f5f Mon Sep 17 00:00:00 2001
From: Yingchi Long <i@lyc.dev>
Date: Thu, 11 Apr 2024 14:17:39 +0800
Subject: [PATCH] libnixt: add attr selecting library (#412)

---
 libnixt/include/nixt/Value.h | 35 +++++++++++++++++++++++
 libnixt/lib/Value.cpp        | 54 ++++++++++++++++++++++++++++++++++++
 libnixt/test/Value.cpp       | 16 +++++++++++
 3 files changed, 105 insertions(+)

diff --git a/libnixt/include/nixt/Value.h b/libnixt/include/nixt/Value.h
index f5be62bc7..423c017da 100644
--- a/libnixt/include/nixt/Value.h
+++ b/libnixt/include/nixt/Value.h
@@ -9,4 +9,39 @@ bool isDerivation(nix::EvalState &State, nix::Value &V);
 std::string attrPathStr(nix::EvalState &State, nix::Value &V,
                         const std::string &AttrPath);
 
+/// \brief Transform a vector of string into a vector of nix symbols.
+std::vector<nix::Symbol> toSymbols(nix::SymbolTable &STable,
+                                   const std::vector<std::string> &Names);
+
+/// \copydoc toSymbols
+std::vector<nix::Symbol> toSymbols(nix::SymbolTable &STable,
+                                   const std::vector<std::string_view> &Names);
+
+/// \brief Select attribute \p Attr
+nix::Value &selectAttr(nix::EvalState &State, nix::Value &V, nix::Symbol Attr);
+
+/// \brief Given an attrpath in nix::Value \p V, select it
+nix::Value &selectAttrPath(nix::EvalState &State, nix::Value &V,
+                           std::vector<nix::Symbol>::const_iterator Begin,
+                           std::vector<nix::Symbol>::const_iterator End);
+
+/// \copydoc selectAttrPath
+inline nix::Value &selectSymbols(nix::EvalState &State, nix::Value &V,
+                                 const std::vector<nix::Symbol> &AttrPath) {
+  return selectAttrPath(State, V, AttrPath.begin(), AttrPath.end());
+}
+
+/// \copydoc selectAttrPath
+inline nix::Value &selectStrings(nix::EvalState &State, nix::Value &V,
+                                 const std::vector<std::string> &AttrPath) {
+  return selectSymbols(State, V, toSymbols(State.symbols, AttrPath));
+}
+
+/// \copydoc selectAttrPath
+inline nix::Value &
+selectStringViews(nix::EvalState &State, nix::Value &V,
+                  const std::vector<std::string_view> &AttrPath) {
+  return selectSymbols(State, V, toSymbols(State.symbols, AttrPath));
+}
+
 } // namespace nixt
diff --git a/libnixt/lib/Value.cpp b/libnixt/lib/Value.cpp
index 3227cab94..cdf03cea8 100644
--- a/libnixt/lib/Value.cpp
+++ b/libnixt/lib/Value.cpp
@@ -1,6 +1,8 @@
 #include "nixt/Value.h"
 
 #include <nix/attr-path.hh>
+#include <nix/nixexpr.hh>
+#include <nix/value.hh>
 
 using namespace nixt;
 
@@ -39,3 +41,55 @@ std::string nixt::attrPathStr(nix::EvalState &State, nix::Value &V,
   State.forceValue(*VPath, Pos);
   return std::string(State.forceStringNoCtx(*VPath, nix::noPos, ""));
 }
+
+std::vector<nix::Symbol>
+nixt::toSymbols(nix::SymbolTable &STable,
+                const std::vector<std::string> &Names) {
+  std::vector<nix::Symbol> Res;
+  Res.reserve(Names.size());
+  for (const auto &Name : Names) {
+    Res.emplace_back(STable.create(Name));
+  }
+  return Res;
+}
+
+std::vector<nix::Symbol>
+nixt::toSymbols(nix::SymbolTable &STable,
+                const std::vector<std::string_view> &Names) {
+  std::vector<nix::Symbol> Res;
+  Res.reserve(Names.size());
+  for (const auto &Name : Names) {
+    Res.emplace_back(STable.create(Name));
+  }
+  return Res;
+}
+
+nix::Value &nixt::selectAttr(nix::EvalState &State, nix::Value &V,
+                             nix::Symbol Attr) {
+  State.forceValue(V, nix::noPos);
+
+  if (V.type() != nix::ValueType::nAttrs)
+    throw nix::TypeError("value is not an attrset");
+
+  assert(V.attrs && "nix must allocate non-null attrs!");
+  auto *Nested = V.attrs->find(Attr);
+  if (Nested == V.attrs->end())
+    throw nix::AttrPathNotFound("attrname " + State.symbols[Attr] +
+                                " not found in attrset");
+
+  assert(Nested->value && "nix must allocate non-null nested value!");
+  return *Nested->value;
+}
+
+/// \brief Given an attrpath in nix::Value \p V, select it
+nix::Value &nixt::selectAttrPath(nix::EvalState &State, nix::Value &V,
+                                 std::vector<nix::Symbol>::const_iterator Begin,
+                                 std::vector<nix::Symbol>::const_iterator End) {
+  // If the attrpath is emtpy, return value itself.
+  if (Begin == End)
+    return V;
+
+  // Otherwise, select it.
+  nix::Value &Nested = selectAttr(State, V, *Begin);
+  return selectAttrPath(State, Nested, ++Begin, End);
+}
diff --git a/libnixt/test/Value.cpp b/libnixt/test/Value.cpp
index 592f97165..c7a1f4710 100644
--- a/libnixt/test/Value.cpp
+++ b/libnixt/test/Value.cpp
@@ -45,4 +45,20 @@ TEST_F(ValueTest, IsDerivation_pos) {
   ASSERT_TRUE(isDerivation(*State, V));
 }
 
+TEST_F(ValueTest, selectAttrPath) {
+  nix::Expr *AST = State->parseExprFromString(R"({ a.b.c.d = 1; })", cwd());
+  nix::Value V;
+  State->eval(AST, V);
+
+  nix::Value &Nested = selectStringViews(*State, V, {"a", "b"});
+
+  // Make sure no extra "force" performed.
+  ASSERT_EQ(Nested.type(), nix::ValueType::nThunk);
+
+  nix::Value &Kern = selectStringViews(*State, Nested, {"c", "d"});
+
+  ASSERT_EQ(Kern.type(), nix::ValueType::nInt);
+  ASSERT_EQ(Kern.integer, 1);
+}
+
 } // namespace