Add libclang API to store and query CAS objects

clang/include/clang-c/CAS.h

@@ -185,6 +185,35 @@ CINDEX_LINKAGE void clang_experimental_cas_loadObjectByString_async(
     void (*Callback)(void *Ctx, CXCASObject, CXError),
     CXCASCancellationToken *OutToken);
 
+/**
+ * Stores an object in the CAS.
+ *
+ * \param Refs The refs of the object to be stored.
+ * \param RefsCount The number of refs.
+ * \param Data The data of the object to be stored.
+ * \param DataSize The size of \Data.
+ * \param[out] OutError The error object to pass back to client (if any).
+ * If non-null the object must be disposed using \c clang_Error_dispose.
+ * \returns The CAS ID of the stored object, or an empty string if an
+ * error occurred. The ID should be disposed using
+ * \c clang_disposeString.
+ */
+CINDEX_LINKAGE CXString clang_experimental_cas_storeObject(
+    CXCASDatabases, const char *const *Refs, size_t RefsCount, const char *Data,
+    size_t DataSize, CXError *OutError);
+
+/**
+ * \returns the CAS object's refs.
+ */
+CINDEX_LINKAGE CXStringSet *
+    clang_experimental_cas_CASObject_getRefs(CXCASObject);
+
+/**
+ * \returns the CAS object's data, which is valid for the lifetime of the
+ * CXCASObject.
+ */
+CINDEX_LINKAGE CXString clang_experimental_cas_CASObject_getData(CXCASObject);
+
 /**
  * Dispose of a \c CXCASObject object.
  */

clang/test/CAS/libclang-cas-api.c

@@ -0,0 +1,19 @@
+// RUN: rm -rf %t && mkdir %t
+
+// RUN: c-index-test core -cas-store -cas-path %t/cas -cas-object-data aGVsbG8sIHdvcmxkIQo= 1> %t/obj1.txt
+
+// RUN: c-index-test core -cas-load -cas-path %t/cas @%t/obj1.txt 1> %t/out1.txt
+// RUN: FileCheck %s -input-file %t/out1.txt -check-prefix=OBJECTONE
+//
+// OBJECTONE: Data:
+// OBJECTONE: aGVsbG8sIHdvcmxkIQo=
+
+// RUN: c-index-test core -cas-store -cas-path %t/cas -cas-object-data Zm9vCg== --cas-object-ref @%t/obj1.txt 1> %t/obj2.txt
+
+// RUN: c-index-test core -cas-load -cas-path %t/cas @%t/obj2.txt 1> %t/out2.txt
+// RUN: FileCheck %s -input-file %t/out2.txt -check-prefix=OBJECTTWO
+//
+// OBJECTTWO: Refs:
+// OBJECTTWO: llvmcas://{{[a-z0-9]+}}
+// OBJECTTWO: Data:
+// OBJECTTWO: Zm9vCg==

clang/tools/c-index-test/core_main.cpp

@@ -9,11 +9,11 @@
 #include "JSONAggregation.h"
 #include "indexstore/IndexStoreCXX.h"
 #include "clang-c/Dependencies.h"
-#include "clang/DirectoryWatcher/DirectoryWatcher.h"
 #include "clang/AST/Mangle.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/PathRemapper.h"
 #include "clang/CodeGen/ObjectFilePCHContainerOperations.h"
+#include "clang/DirectoryWatcher/DirectoryWatcher.h"
 #include "clang/Frontend/ASTUnit.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
@@ -29,6 +29,7 @@
 #include "clang/Serialization/ASTReader.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/Support/Base64.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
@@ -61,6 +62,8 @@ enum class ActionType {
   UploadCachedJob,
   MaterializeCachedJob,
   ReplayCachedJob,
+  CASStore,
+  CASLoad,
   PruneCAS,
   WatchDir,
 };
@@ -92,6 +95,10 @@ Action(cl::desc("Action:"), cl::init(ActionType::None),
                      "Materialize cached compilation data from upstream CAS"),
           clEnumValN(ActionType::ReplayCachedJob, "replay-cached-job",
                      "Replay a cached compilation from the CAS"),
+          clEnumValN(ActionType::CASStore, "cas-store",
+                     "Store an object in the CAS"),
+          clEnumValN(ActionType::CASLoad, "cas-load",
+                     "Load an object from the CAS"),
           clEnumValN(ActionType::PruneCAS, "prune-cas", "Prune CAS data"),
           clEnumValN(ActionType::WatchDir,
                      "watch-dir", "Watch directory for file events")),
@@ -155,6 +162,11 @@ static llvm::cl::opt<std::string>
     CASPluginPath("fcas-plugin-path", llvm::cl::desc("Path for CAS plugin"));
 static cl::list<std::string> CASPluginOpts("fcas-plugin-option",
                                            cl::desc("Plugin CAS Options"));
+static cl::list<std::string> CASObjectRefs("cas-object-ref",
+                                           cl::desc("CAS object ref"));
+static llvm::cl::opt<std::string>
+    CASObjectData("cas-object-data",
+                  llvm::cl::desc("CAS object data (base 64)"));
 static llvm::cl::opt<std::string>
     WorkingDir("working-dir", llvm::cl::desc("Path for working directory"));
 static cl::opt<bool> TestCASCancellation(
@@ -1096,6 +1108,58 @@ static int replayCachedJob(ArrayRef<const char *> Args,
   return 0;
 }
 
+static int casLoad(std::string CASID, CXCASDatabases DBs) {
+  CXError Err = nullptr;
+  CXCASObject Object =
+      clang_experimental_cas_loadObjectByString(DBs, CASID.c_str(), &Err);
+  if (Err) {
+    llvm::errs() << clang_Error_getDescription(Err) << "\n";
+    clang_Error_dispose(Err);
+    return 1;
+  }
+  CXStringSet *Refs = clang_experimental_cas_CASObject_getRefs(Object);
+  CXString Data = clang_experimental_cas_CASObject_getData(Object);
+  if (Refs->Count > 0) {
+    llvm::outs() << "Refs:\n";
+    for (unsigned i = 0; i < Refs->Count; ++i) {
+      llvm::outs() << clang_getCString(Refs->Strings[i]) << "\n";
+    }
+  }
+  llvm::outs() << "Data: \n"
+               << encodeBase64(StringRef(clang_getCString(Data))) << "\n";
+  clang_experimental_cas_CASObject_dispose(Object);
+  clang_disposeStringSet(Refs);
+  return 0;
+}
+
+static int casStore(ArrayRef<std::string> Refs, std::string Base64Data,
+                    CXCASDatabases DBs) {
+  std::vector<char *> CRefs;
+  for (std::string Ref : Refs) {
+    CRefs.push_back(strdup(Ref.c_str()));
+  }
+  std::vector<char> Data;
+  if (Error Err = decodeBase64(Base64Data, Data)) {
+    llvm::errs() << Err << "\n";
+    return 1;
+  }
+  CXError Err = nullptr;
+  CXString CASID = clang_experimental_cas_storeObject(
+      DBs, CRefs.data(), CRefs.size(), Data.data(), Data.size(), &Err);
+  for (char *CRef : CRefs) {
+    free(CRef);
+  }
+  if (Err) {
+    clang_disposeString(CASID);
+    llvm::errs() << clang_Error_getDescription(Err) << "\n";
+    clang_Error_dispose(Err);
+    return 1;
+  }
+  llvm::outs() << clang_getCString(CASID) << "\n";
+  clang_disposeString(CASID);
+  return 0;
+}
+
 static int pruneCAS(int64_t Limit, CXCASDatabases DBs) {
   CXError Err = nullptr;
   int64_t Size = clang_experimental_cas_Databases_get_storage_size(DBs, &Err);
@@ -1534,6 +1598,26 @@ int indextest_core_main(int argc, const char **argv) {
                            options::InputFiles[0], DBs);
   }
 
+  if (options::Action == ActionType::CASStore) {
+    if (!DBs) {
+      errs() << "error: CAS was not configured\n";
+      return 1;
+    }
+    return casStore(options::CASObjectRefs, options::CASObjectData, DBs);
+  }
+
+  if (options::Action == ActionType::CASLoad) {
+    if (!DBs) {
+      errs() << "error: CAS was not configured\n";
+      return 1;
+    }
+    if (options::InputFiles.size() != 1) {
+      errs() << "error: expected a single CAS ID as input\n";
+      return 1;
+    }
+    return casLoad(options::InputFiles[0], DBs);
+  }
+
   if (options::Action == ActionType::PruneCAS) {
     if (options::InputFiles.empty()) {
       errs() << "error: missing size limit\n";

clang/tools/libclang/CCAS.cpp

@@ -397,6 +397,51 @@ void clang_experimental_cas_loadObjectByString_async(
   WL->visit(*Ref, /*IsRootNode*/ true);
 }
 
+CXString clang_experimental_cas_storeObject(CXCASDatabases CDBs,
+                                            const char *const *IDStrings,
+                                            size_t IDStringsCount,
+                                            const char *Data, size_t DataSize,
+                                            CXError *OutError) {
+  WrappedCASDatabases &DBs = *unwrap(CDBs);
+  ObjectStore &CAS = *DBs.CAS;
+
+  SmallVector<ObjectRef, 1> Refs;
+  for (size_t i = 0; i < IDStringsCount; ++i) {
+    Expected<CASID> ID = CAS.parseID(IDStrings[i]);
+    if (!ID) {
+      *OutError = cxerror::create(ID.takeError());
+      return cxstring::createEmpty();
+    }
+    std::optional<ObjectRef> Ref = CAS.getReference(*ID);
+    if (!Ref) {
+      *OutError =
+          cxerror::create("Could not get reference for ID: " + ID->toString());
+      return cxstring::createEmpty();
+    }
+    Refs.push_back(*Ref);
+  }
+
+  Expected<ObjectRef> Ref = CAS.store(Refs, ArrayRef(Data, DataSize));
+  if (!Ref) {
+    *OutError = cxerror::create(Ref.takeError());
+    return cxstring::createEmpty();
+  }
+
+  return cxstring::createDup(CAS.getID(*Ref).toString());
+}
+
+CXStringSet *clang_experimental_cas_CASObject_getRefs(CXCASObject CObj) {
+  std::vector<std::string> Refs;
+  for (size_t i = 0; i < unwrap(CObj)->Obj.getNumReferences(); ++i) {
+    Refs.push_back(unwrap(CObj)->Obj.getReferenceID(i).toString());
+  }
+  return cxstring::createSet(Refs);
+}
+
+CXString clang_experimental_cas_CASObject_getData(CXCASObject CObj) {
+  return cxstring::createRef(unwrap(CObj)->Obj.getData());
+}
+
 void clang_experimental_cas_CASObject_dispose(CXCASObject CObj) {
   delete unwrap(CObj);
 }

clang/tools/libclang/libclang.map

@@ -484,6 +484,8 @@ LLVM_16 {
     clang_experimental_cas_CancellationToken_cancel;
     clang_experimental_cas_CancellationToken_dispose;
     clang_experimental_cas_CASObject_dispose;
+    clang_experimental_cas_CASObject_getData;
+    clang_experimental_cas_CASObject_getRefs;
     clang_experimental_cas_Databases_create;
     clang_experimental_cas_Databases_dispose;
     clang_experimental_cas_Databases_get_storage_size;
@@ -504,6 +506,7 @@ LLVM_16 {
     clang_experimental_cas_replayCompilation;
     clang_experimental_cas_ReplayResult_dispose;
     clang_experimental_cas_ReplayResult_getStderr;
+    clang_experimental_cas_storeObject;
     clang_experimental_DependencyScannerService_create_v1;
     clang_experimental_DependencyScannerServiceOptions_create;
     clang_experimental_DependencyScannerServiceOptions_dispose;