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