saja 0.1.0

/*
 * Clang plugin for the #pragma use directive.
 *
 * Copyright (C) 2026  Madeleine Choi
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

#include "clang/Basic/Diagnostic.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Frontend/CompilerInstance.h"
#include "clang/Lex/HeaderSearch.h"
#include "clang/Lex/Preprocessor.h"
#include "clang/Lex/Pragma.h"
#include "clang/Lex/Token.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/Support/Path.h"

#include <string>
#include <vector>

using namespace clang;

namespace {

class UsePragmaHandler final : public PragmaHandler {
public:
  UsePragmaHandler() : PragmaHandler("use") {}

  void HandlePragma(Preprocessor &PP, PragmaIntroducer Introducer,
                    Token &FirstToken) override {
    std::vector<std::string> Segments;
    Token Tok = FirstToken;

    if (Tok.is(tok::identifier) && PP.getSpelling(Tok) == "use")
      PP.Lex(Tok);

    if (!parsePath(PP, Tok, Segments))
      return;

    consumeTree(PP, Tok);

    if (Tok.isNot(tok::eod)) {
      report(PP, Tok.getLocation(), "expected end of #pragma use directive");
      discardLine(PP, Tok);
      return;
    }

    includeBestMatch(PP, Introducer.Loc, Segments);
  }

private:
  static unsigned customDiagnostic(Preprocessor &PP) {
    return PP.getDiagnostics().getCustomDiagID(DiagnosticsEngine::Error, "%0");
  }

  static void report(Preprocessor &PP, SourceLocation Loc, const char *Message) {
    PP.getDiagnostics().Report(Loc, customDiagnostic(PP)) << Message;
  }

  static void discardLine(Preprocessor &PP, Token &Tok) {
    while (Tok.isNot(tok::eod))
      PP.Lex(Tok);
  }

  static bool consumePathSeparator(Preprocessor &PP, Token &Tok) {
    if (Tok.is(tok::coloncolon)) {
      PP.Lex(Tok);
      return true;
    }

    if (Tok.isNot(tok::colon))
      return false;

    Token SecondColon;
    PP.Lex(SecondColon);
    if (SecondColon.isNot(tok::colon)) {
      report(PP, SecondColon.getLocation(),
             "expected ':' after ':' in #pragma use path");
      Tok = SecondColon;
      return false;
    }

    PP.Lex(Tok);
    return true;
  }

  static bool parsePath(Preprocessor &PP, Token &Tok,
                        std::vector<std::string> &Segments) {
    if (Tok.isNot(tok::identifier)) {
      report(PP, Tok.getLocation(), "expected path after #pragma use");
      discardLine(PP, Tok);
      return false;
    }

    while (true) {
      if (Tok.isNot(tok::identifier)) {
        report(PP, Tok.getLocation(), "expected identifier in #pragma use path");
        discardLine(PP, Tok);
        return false;
      }

      Segments.push_back(PP.getSpelling(Tok));
      PP.Lex(Tok);

      if (!consumePathSeparator(PP, Tok))
        return true;

      if (Tok.is(tok::star) || Tok.is(tok::l_brace))
        return true;
    }
  }

  static void consumeTree(Preprocessor &PP, Token &Tok) {
    if (Tok.is(tok::star)) {
      PP.Lex(Tok);
      return;
    }

    if (Tok.isNot(tok::l_brace))
      return;

    unsigned Depth = 1;
    while (Depth > 0) {
      PP.Lex(Tok);

      if (Tok.is(tok::eod))
        return;
      if (Tok.is(tok::l_brace))
        ++Depth;
      else if (Tok.is(tok::r_brace))
        --Depth;
    }

    PP.Lex(Tok);
  }

  static std::string headerName(ArrayRef<std::string> Segments) {
    llvm::SmallString<128> Path;

    for (const std::string &Segment : Segments)
      llvm::sys::path::append(Path, Segment);

    llvm::sys::path::replace_extension(Path, "h");
    return std::string(Path);
  }

  static void includeBestMatch(Preprocessor &PP, SourceLocation Loc,
                               std::vector<std::string> Segments) {
    while (!Segments.empty()) {
      std::string Header = headerName(Segments);

      ConstSearchDirIterator CurDir = nullptr;
      llvm::SmallString<128> SearchPath;
      llvm::SmallString<128> RelativePath;
      ModuleMap::KnownHeader SuggestedModule;
      bool IsMapped = false;
      bool IsFrameworkFound = false;

      OptionalFileEntryRef File = PP.LookupFile(
          Loc, Header, false, nullptr, nullptr, &CurDir, &SearchPath,
          &RelativePath, &SuggestedModule, &IsMapped, &IsFrameworkFound);

      if (File) {
        SourceManager &SM = PP.getSourceManager();
        SrcMgr::CharacteristicKind FileKind =
            PP.getHeaderSearchInfo().getFileDirFlavor(*File);
        FileID FID = SM.createFileID(*File, Loc, FileKind);
        PP.EnterSourceFile(FID, CurDir, Loc);
        return;
      }

      Segments.pop_back();
    }

    report(PP, Loc, "could not resolve #pragma use path to a generated header; are there any exports?");
  }
};

} // namespace

static PragmaHandlerRegistry::Add<UsePragmaHandler>
    X("use", "expand Saja #pragma use directives");