path: root/Driver/RewriteTest.cpp

//===--- RewriteTest.cpp - Playground for the code rewriter ---------------===//
//
//                     The LLVM Compiler Infrastructure
//
// This file was developed by Chris Lattner and is distributed under the
// University of Illinois Open Source License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// Hacks and fun related to the code rewriter.
//
//===----------------------------------------------------------------------===//

#include "ASTConsumers.h"
#include "clang/Rewrite/Rewriter.h"
#include "clang/AST/AST.h"
#include "clang/AST/ASTConsumer.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Basic/IdentifierTable.h"
#include "llvm/ADT/StringExtras.h"
using namespace clang;
using llvm::utostr;

namespace {
  class RewriteTest : public ASTConsumer {
    Rewriter Rewrite;
    ASTContext *Context;
    SourceManager *SM;
    unsigned MainFileID;
    SourceLocation LastIncLoc;
    llvm::SmallVector<ObjcImplementationDecl *, 8> ClassImplementation;
    llvm::SmallVector<ObjcCategoryImplDecl *, 8> CategoryImplementation;
    
    FunctionDecl *MsgSendFunctionDecl;
    FunctionDecl *GetClassFunctionDecl;
    FunctionDecl *SelGetUidFunctionDecl;
    
    static const int OBJC_ABI_VERSION =7 ;
  public:
    void Initialize(ASTContext &context, unsigned mainFileID) {
      Context = &context;
      SM = &Context->SourceMgr;
      MainFileID = mainFileID;
      MsgSendFunctionDecl = 0;
      GetClassFunctionDecl = 0;
      SelGetUidFunctionDecl = 0;
      Rewrite.setSourceMgr(Context->SourceMgr);
    }

    // Top Level Driver code.
    virtual void HandleTopLevelDecl(Decl *D);
    void HandleDeclInMainFile(Decl *D);
    ~RewriteTest();

    // Syntactic Rewriting.
    void RewriteInclude(SourceLocation Loc);
    void RewriteTabs();
    void RewriteForwardClassDecl(ObjcClassDecl *Dcl);
    
    // Expression Rewriting.
    Stmt *RewriteFunctionBody(Stmt *S);
    Stmt *RewriteAtEncode(ObjCEncodeExpr *Exp);
    Stmt *RewriteMessageExpr(ObjCMessageExpr *Exp);
    CallExpr *SynthesizeCallToFunctionDecl(FunctionDecl *FD, 
                                           Expr **args, unsigned nargs);
    // Metadata emission.
    void RewriteObjcClassMetaData(ObjcImplementationDecl *IDecl,
                                  std::string &Result);
    
    void RewriteObjcCategoryImplDecl(ObjcCategoryImplDecl *CDecl,
                                     std::string &Result);
    
    void RewriteObjcMethodsMetaData(ObjcMethodDecl **Methods,
                                    int NumMethods,
                                    bool IsInstanceMethod,
                                    const char *prefix,
                                    const char *ClassName,
                                    std::string &Result);
    
    void RewriteObjcProtocolsMetaData(ObjcProtocolDecl **Protocols,
                                      int NumProtocols,
                                      const char *prefix,
                                      const char *ClassName,
                                      std::string &Result);
    void WriteObjcMetaData(std::string &Result);
  };
}

ASTConsumer *clang::CreateCodeRewriterTest() { return new RewriteTest(); }

//===----------------------------------------------------------------------===//
// Top Level Driver Code
//===----------------------------------------------------------------------===//

void RewriteTest::HandleTopLevelDecl(Decl *D) {
  // Two cases: either the decl could be in the main file, or it could be in a
  // #included file.  If the former, rewrite it now.  If the later, check to see
  // if we rewrote the #include/#import.
  SourceLocation Loc = D->getLocation();
  Loc = SM->getLogicalLoc(Loc);
  
  // If this is for a builtin, ignore it.
  if (Loc.isInvalid()) return;

  // Look for built-in declarations that we need to refer during the rewrite.
  if (FunctionDecl *FD = dyn_cast<FunctionDecl>(D)) {
    if (strcmp(FD->getName(), "objc_msgSend") == 0)
      MsgSendFunctionDecl = FD;
    else if (strcmp(FD->getName(), "objc_getClass") == 0)
      GetClassFunctionDecl = FD;
    else if (strcmp(FD->getName(), "sel_getUid") == 0)
      SelGetUidFunctionDecl = FD;
  }
  
  // If we have a decl in the main file, see if we should rewrite it.
  if (SM->getDecomposedFileLoc(Loc).first == MainFileID)
    return HandleDeclInMainFile(D);

  // Otherwise, see if there is a #import in the main file that should be
  // rewritten.
  RewriteInclude(Loc);
}

/// HandleDeclInMainFile - This is called for each top-level decl defined in the
/// main file of the input.
void RewriteTest::HandleDeclInMainFile(Decl *D) {
  if (FunctionDecl *FD = dyn_cast<FunctionDecl>(D))
    if (Stmt *Body = FD->getBody())
      FD->setBody(RewriteFunctionBody(Body));
  
  if (ObjcImplementationDecl *CI = dyn_cast<ObjcImplementationDecl>(D))
    ClassImplementation.push_back(CI);
  else if (ObjcCategoryImplDecl *CI = dyn_cast<ObjcCategoryImplDecl>(D))
    CategoryImplementation.push_back(CI);
  else if (ObjcClassDecl