Home | History | Annotate | Download | only in Rewrite 
      1 //===--- Rewriter.cpp - Code rewriting interface --------------------------===//
      2 //
      3 //                     The LLVM Compiler Infrastructure
      4 //
      5 // This file is distributed under the University of Illinois Open Source
      6 // License. See LICENSE.TXT for details.
      7 //
      8 //===----------------------------------------------------------------------===//
      9 //
     10 //  This file defines the Rewriter class, which is used for code
     11 //  transformations.
     12 //
     13 //===----------------------------------------------------------------------===//
     14 
     15 #include "clang/Rewrite/Rewriter.h"
     16 #include "clang/AST/Stmt.h"
     17 #include "clang/AST/Decl.h"
     18 #include "clang/Lex/Lexer.h"
     19 #include "clang/Basic/SourceManager.h"
     20 #include "llvm/ADT/SmallString.h"
     21 using namespace clang;
     22 
     23 raw_ostream &RewriteBuffer::write(raw_ostream &os) const {
     24   // FIXME: eliminate the copy by writing out each chunk at a time
     25   os << std::string(begin(), end());
     26   return os;
     27 }
     28 
     29 /// \brief Return true if this character is non-new-line whitespace:
     30 /// ' ', '\t', '\f', '\v', '\r'.
     31 static inline bool isWhitespace(unsigned char c) {
     32   switch (c) {
     33   case ' ':
     34   case '\t':
     35   case '\f':
     36   case '\v':
     37   case '\r':
     38     return true;
     39   default:
     40     return false;
     41   }
     42 }
     43 
     44 void RewriteBuffer::RemoveText(unsigned OrigOffset, unsigned Size,
     45                                bool removeLineIfEmpty) {
     46   // Nothing to remove, exit early.
     47   if (Size == 0) return;
     48 
     49   unsigned RealOffset = getMappedOffset(OrigOffset, true);
     50   assert(RealOffset+Size < Buffer.size() && "Invalid location");
     51 
     52   // Remove the dead characters.
     53   Buffer.erase(RealOffset, Size);
     54 
     55   // Add a delta so that future changes are offset correctly.
     56   AddReplaceDelta(OrigOffset, -Size);
     57 
     58   if (removeLineIfEmpty) {
     59     // Find the line that the remove occurred and if it is completely empty
     60     // remove the line as well.
     61 
     62     iterator curLineStart = begin();
     63     unsigned curLineStartOffs = 0;
     64     iterator posI = begin();
     65     for (unsigned i = 0; i != RealOffset; ++i) {
     66       if (*posI == '\n') {
     67         curLineStart = posI;
     68         ++curLineStart;
     69         curLineStartOffs = i + 1;
     70       }
     71       ++posI;
     72     }
     73 
     74     unsigned lineSize = 0;
     75     posI = curLineStart;
     76     while (posI != end() && isWhitespace(*posI)) {
     77       ++posI;
     78       ++lineSize;
     79     }
     80     if (posI != end() && *posI == '\n') {
     81       Buffer.erase(curLineStartOffs, lineSize + 1/* + '\n'*/);
     82       AddReplaceDelta(curLineStartOffs, -(lineSize + 1/* + '\n'*/));
     83     }
     84   }
     85 }
     86 
     87 void RewriteBuffer::InsertText(unsigned OrigOffset, StringRef Str,
     88                                bool InsertAfter) {
     89 
     90   // Nothing to insert, exit early.
     91   if (Str.empty()) return;
     92 
     93   unsigned RealOffset = getMappedOffset(OrigOffset, InsertAfter);
     94   Buffer.insert(RealOffset, Str.begin(), Str.end());
     95 
     96   // Add a delta so that future changes are offset correctly.
     97   AddInsertDelta(OrigOffset, Str.size());
     98 }
     99 
    100 /// ReplaceText - This method replaces a range of characters in the input
    101 /// buffer with a new string.  This is effectively a combined "remove+insert"
    102 /// operation.
    103 void RewriteBuffer::ReplaceText(unsigned OrigOffset, unsigned OrigLength,
    104                                 StringRef NewStr) {
    105   unsigned RealOffset = getMappedOffset(OrigOffset, true);
    106   Buffer.erase(RealOffset, OrigLength);
    107   Buffer.insert(RealOffset, NewStr.begin(), NewStr.end());
    108   if (OrigLength != NewStr.size())
    109     AddReplaceDelta(OrigOffset, NewStr.size() - OrigLength);
    110 }
    111 
    112 
    113 //===----------------------------------------------------------------------===//
    114 // Rewriter class
    115 //===----------------------------------------------------------------------===//
    116 
    117 /// getRangeSize - Return the size in bytes of the specified range if they
    118 /// are in the same file.  If not, this returns -1.
    119 int Rewriter::getRangeSize(const CharSourceRange &Range,
    120                            RewriteOptions opts) const {
    121   if (!isRewritable(Range.getBegin()) ||
    122       !isRewritable(Range.getEnd())) return -1;
    123 
    124   FileID StartFileID, EndFileID;
    125   unsigned StartOff, EndOff;
    126 
    127   StartOff = getLocationOffsetAndFileID(Range.getBegin(), StartFileID);
    128   EndOff   = getLocationOffsetAndFileID(Range.getEnd(), EndFileID);
    129 
    130   if (StartFileID != EndFileID)
    131     return -1;
    132 
    133   // If edits have been made to this buffer, the delta between the range may
    134   // have changed.
    135   std::map<FileID, RewriteBuffer>::const_iterator I =
    136     RewriteBuffers.find(StartFileID);
    137   if (I != RewriteBuffers.end()) {
    138     const RewriteBuffer &RB = I->second;
    139     EndOff = RB.getMappedOffset(EndOff, opts.IncludeInsertsAtEndOfRange);
    140     StartOff = RB.getMappedOffset(StartOff, !opts.IncludeInsertsAtBeginOfRange);
    141   }
    142 
    143 
    144   // Adjust the end offset to the end of the last token, instead of being the
    145   // start of the last token if this is a token range.
    146   if (Range.isTokenRange())
    147     EndOff += Lexer::MeasureTokenLength(Range.getEnd(), *SourceMgr, *LangOpts);
    148 
    149   return EndOff-StartOff;
    150 }
    151 
    152 int Rewriter::getRangeSize(SourceRange Range, RewriteOptions opts) const {
    153   return getRangeSize(CharSourceRange::getTokenRange(Range), opts);
    154 }
    155 
    156 
    157 /// getRewrittenText - Return the rewritten form of the text in the specified
    158 /// range.  If the start or end of the range was unrewritable or if they are
    159 /// in different buffers, this returns an empty string.
    160 ///
    161 /// Note that this method is not particularly efficient.
    162 ///
    163 std::string Rewriter::getRewrittenText(SourceRange Range) const {
    164   if (!isRewritable(Range.getBegin()) ||
    165       !isRewritable(Range.getEnd()))
    166     return "";
    167 
    168   FileID StartFileID, EndFileID;
    169   unsigned StartOff, EndOff;
    170   StartOff = getLocationOffsetAndFileID(Range.getBegin(), StartFileID);
    171   EndOff   = getLocationOffsetAndFileID(Range.getEnd(), EndFileID);
    172 
    173   if (StartFileID != EndFileID)
    174     return ""; // Start and end in different buffers.
    175 
    176   // If edits have been made to this buffer, the delta between the range may
    177   // have changed.
    178   std::map<FileID, RewriteBuffer>::const_iterator I =
    179     RewriteBuffers.find(StartFileID);
    180   if (I == RewriteBuffers.end()) {
    181     // If the buffer hasn't been rewritten, just return the text from the input.
    182     const char *Ptr = SourceMgr->getCharacterData(Range.getBegin());
    183 
    184     // Adjust the end offset to the end of the last token, instead of being the
    185     // start of the last token.
    186     EndOff += Lexer::MeasureTokenLength(Range.getEnd(), *SourceMgr, *LangOpts);
    187     return std::string(Ptr, Ptr+EndOff-StartOff);
    188   }
    189 
    190   const RewriteBuffer &RB = I->second;
    191   EndOff = RB.getMappedOffset(EndOff, true);
    192   StartOff = RB.getMappedOffset(StartOff);
    193 
    194   // Adjust the end offset to the end of the last token, instead of being the
    195   // start of the last token.
    196   EndOff += Lexer::MeasureTokenLength(Range.getEnd(), *SourceMgr, *LangOpts);
    197 
    198   // Advance the iterators to the right spot, yay for linear time algorithms.
    199   RewriteBuffer::iterator Start = RB.begin();
    200   std::advance(Start, StartOff);
    201   RewriteBuffer::iterator End = Start;
    202   std::advance(End, EndOff-StartOff);
    203 
    204   return std::string(Start, End);
    205 }
    206 
    207 unsigned Rewriter::getLocationOffsetAndFileID(SourceLocation Loc,
    208                                               FileID &FID) const {
    209   assert(Loc.isValid() && "Invalid location");
    210   std::pair<FileID,unsigned> V = SourceMgr->getDecomposedLoc(Loc);
    211   FID = V.first;
    212   return V.second;
    213 }
    214 
    215 
    216 /// getEditBuffer - Get or create a RewriteBuffer for the specified FileID.
    217 ///
    218 RewriteBuffer &Rewriter::getEditBuffer(FileID FID) {
    219   std::map<FileID, RewriteBuffer>::iterator I =
    220     RewriteBuffers.lower_bound(FID);
    221   if (I != RewriteBuffers.end() && I->first == FID)
    222     return I->second;
    223   I = RewriteBuffers.insert(I, std::make_pair(FID, RewriteBuffer()));
    224 
    225   StringRef MB = SourceMgr->getBufferData(FID);
    226   I->second.Initialize(MB.begin(), MB.end());
    227 
    228   return I->second;
    229 }
    230 
    231 /// InsertText - Insert the specified string at the specified location in the
    232 /// original buffer.
    233 bool Rewriter::InsertText(SourceLocation Loc, StringRef Str,
    234                           bool InsertAfter, bool indentNewLines) {
    235   if (!isRewritable(Loc)) return true;
    236   FileID FID;
    237   unsigned StartOffs = getLocationOffsetAndFileID(Loc, FID);
    238 
    239   SmallString<128> indentedStr;
    240   if (indentNewLines && Str.find('\n') != StringRef::npos) {
    241     StringRef MB = SourceMgr->getBufferData(FID);
    242 
    243     unsigned lineNo = SourceMgr->getLineNumber(FID, StartOffs) - 1;
    244     const SrcMgr::ContentCache *
    245         Content = SourceMgr->getSLocEntry(FID).getFile().getContentCache();
    246     unsigned lineOffs = Content->SourceLineCache[lineNo];
    247 
    248     // Find the whitespace at the start of the line.
    249     StringRef indentSpace;
    250     {
    251       unsigned i = lineOffs;
    252       while (isWhitespace(MB[i]))
    253         ++i;
    254       indentSpace = MB.substr(lineOffs, i-lineOffs);
    255     }
    256 
    257     SmallVector<StringRef, 4> lines;
    258     Str.split(lines, "\n");
    259 
    260     for (unsigned i = 0, e = lines.size(); i != e; ++i) {
    261       indentedStr += lines[i];
    262       if (i < e-1) {
    263         indentedStr += '\n';
    264         indentedStr += indentSpace;
    265       }
    266     }
    267     Str = indentedStr.str();
    268   }
    269 
    270   getEditBuffer(FID).InsertText(StartOffs, Str, InsertAfter);
    271   return false;
    272 }
    273 
    274 bool Rewriter::InsertTextAfterToken(SourceLocation Loc, StringRef Str) {
    275   if (!isRewritable(Loc)) return true;
    276   FileID FID;
    277   unsigned StartOffs = getLocationOffsetAndFileID(Loc, FID);
    278   RewriteOptions rangeOpts;
    279   rangeOpts.IncludeInsertsAtBeginOfRange = false;
    280   StartOffs += getRangeSize(SourceRange(Loc, Loc), rangeOpts);
    281   getEditBuffer(FID).InsertText(StartOffs, Str, /*InsertAfter*/true);
    282   return false;
    283 }
    284 
    285 /// RemoveText - Remove the specified text region.
    286 bool Rewriter::RemoveText(SourceLocation Start, unsigned Length,
    287                           RewriteOptions opts) {
    288   if (!isRewritable(Start)) return true;
    289   FileID FID;
    290   unsigned StartOffs = getLocationOffsetAndFileID(Start, FID);
    291   getEditBuffer(FID).RemoveText(StartOffs, Length, opts.RemoveLineIfEmpty);
    292   return false;
    293 }
    294 
    295 /// ReplaceText - This method replaces a range of characters in the input
    296 /// buffer with a new string.  This is effectively a combined "remove/insert"
    297 /// operation.
    298 bool Rewriter::ReplaceText(SourceLocation Start, unsigned OrigLength,
    299                            StringRef NewStr) {
    300   if (!isRewritable(Start)) return true;
    301   FileID StartFileID;
    302   unsigned StartOffs = getLocationOffsetAndFileID(Start, StartFileID);
    303 
    304   getEditBuffer(StartFileID).ReplaceText(StartOffs, OrigLength, NewStr);
    305   return false;
    306 }
    307 
    308 bool Rewriter::ReplaceText(SourceRange range, SourceRange replacementRange) {
    309   if (!isRewritable(range.getBegin())) return true;
    310   if (!isRewritable(range.getEnd())) return true;
    311   if (replacementRange.isInvalid()) return true;
    312   SourceLocation start = range.getBegin();
    313   unsigned origLength = getRangeSize(range);
    314   unsigned newLength = getRangeSize(replacementRange);
    315   FileID FID;
    316   unsigned newOffs = getLocationOffsetAndFileID(replacementRange.getBegin(),
    317                                                 FID);
    318   StringRef MB = SourceMgr->getBufferData(FID);
    319   return ReplaceText(start, origLength, MB.substr(newOffs, newLength));
    320 }
    321 
    322 /// ReplaceStmt - This replaces a Stmt/Expr with another, using the pretty
    323 /// printer to generate the replacement code.  This returns true if the input
    324 /// could not be rewritten, or false if successful.
    325 bool Rewriter::ReplaceStmt(Stmt *From, Stmt *To) {
    326   // Measaure the old text.
    327   int Size = getRangeSize(From->getSourceRange());
    328   if (Size == -1)
    329     return true;
    330 
    331   // Get the new text.
    332   std::string SStr;
    333   llvm::raw_string_ostream S(SStr);
    334   To->printPretty(S, 0, PrintingPolicy(*LangOpts));
    335   const std::string &Str = S.str();
    336 
    337   ReplaceText(From->getLocStart(), Size, Str);
    338   return false;
    339 }
    340 
    341 std::string Rewriter::ConvertToString(Stmt *From) {
    342   std::string SStr;
    343   llvm::raw_string_ostream S(SStr);
    344   From->printPretty(S, 0, PrintingPolicy(*LangOpts));
    345   return S.str();
    346 }
    347 
    348 bool Rewriter::IncreaseIndentation(CharSourceRange range,
    349                                    SourceLocation parentIndent) {
    350   if (range.isInvalid()) return true;
    351   if (!isRewritable(range.getBegin())) return true;
    352   if (!isRewritable(range.getEnd())) return true;
    353   if (!isRewritable(parentIndent)) return true;
    354 
    355   FileID StartFileID, EndFileID, parentFileID;
    356   unsigned StartOff, EndOff, parentOff;
    357 
    358   StartOff = getLocationOffsetAndFileID(range.getBegin(), StartFileID);
    359   EndOff   = getLocationOffsetAndFileID(range.getEnd(), EndFileID);
    360   parentOff = getLocationOffsetAndFileID(parentIndent, parentFileID);
    361 
    362   if (StartFileID != EndFileID || StartFileID != parentFileID)
    363     return true;
    364   if (StartOff > EndOff)
    365     return true;
    366 
    367   FileID FID = StartFileID;
    368   StringRef MB = SourceMgr->getBufferData(FID);
    369 
    370   unsigned parentLineNo = SourceMgr->getLineNumber(FID, parentOff) - 1;
    371   unsigned startLineNo = SourceMgr->getLineNumber(FID, StartOff) - 1;
    372   unsigned endLineNo = SourceMgr->getLineNumber(FID, EndOff) - 1;
    373 
    374   const SrcMgr::ContentCache *
    375       Content = SourceMgr->getSLocEntry(FID).getFile().getContentCache();
    376 
    377   // Find where the lines start.
    378   unsigned parentLineOffs = Content->SourceLineCache[parentLineNo];
    379   unsigned startLineOffs = Content->SourceLineCache[startLineNo];
    380 
    381   // Find the whitespace at the start of each line.
    382   StringRef parentSpace, startSpace;
    383   {
    384     unsigned i = parentLineOffs;
    385     while (isWhitespace(MB[i]))
    386       ++i;
    387     parentSpace = MB.substr(parentLineOffs, i-parentLineOffs);
    388 
    389     i = startLineOffs;
    390     while (isWhitespace(MB[i]))
    391       ++i;
    392     startSpace = MB.substr(startLineOffs, i-startLineOffs);
    393   }
    394   if (parentSpace.size() >= startSpace.size())
    395     return true;
    396   if (!startSpace.startswith(parentSpace))
    397     return true;
    398 
    399   StringRef indent = startSpace.substr(parentSpace.size());
    400 
    401   // Indent the lines between start/end offsets.
    402   RewriteBuffer &RB = getEditBuffer(FID);
    403   for (unsigned lineNo = startLineNo; lineNo <= endLineNo; ++lineNo) {
    404     unsigned offs = Content->SourceLineCache[lineNo];
    405     unsigned i = offs;
    406     while (isWhitespace(MB[i]))
    407       ++i;
    408     StringRef origIndent = MB.substr(offs, i-offs);
    409     if (origIndent.startswith(startSpace))
    410       RB.InsertText(offs, indent, /*InsertAfter=*/false);
    411   }
    412 
    413   return false;
    414 }
    415