Playing with the AST rewriter.

tests/manual/cplusplus/main.cpp

@@ -32,40 +32,172 @@
 ***************************************************************************/
 
 #include <AST.h>
+#include <ASTVisitor.h>
 #include <Control.h>
 #include <Scope.h>
 #include <Semantic.h>
 #include <TranslationUnit.h>
 
 #include <QtCore/QFile>
+#include <QtCore/QList>
+#include <QtDebug>
 
 #include <cstdio>
 #include <cstdlib>
 
+class Rewrite
+{
+    QMultiMap<unsigned, QByteArray> _insertBefore;
+    QMultiMap<unsigned, QByteArray> _insertAfter;
+    QSet<unsigned> _removed;
+
+public:
+    void remove(unsigned index)
+    { remove(index, index + 1); }
+
+    void remove(unsigned first, unsigned last)
+    {
+        Q_ASSERT(first < last);
+
+        for (; first != last; ++first)
+            _removed.insert(first);
+    }
+
+    void insertTextBefore(unsigned index, const QByteArray &text)
+    { _insertBefore.insert(index, text); }
+
+    void insertTextAfter(unsigned index, const QByteArray &text)
+    { _insertAfter.insert(index, text); }
+
+    void rewrite(const TranslationUnit *unit,
+                 const QByteArray &contents,
+                 QByteArray *out) const
+    {
+        const char *source = contents.constData();
+        unsigned previousTokenEndPosition = 0;
+        for (unsigned i = 0; i < unit->tokenCount(); ++i) {
+            const Token &tk = unit->tokenAt(i);
+
+            if (previousTokenEndPosition != tk.begin()) {
+                Q_ASSERT(previousTokenEndPosition < tk.begin());
+                out->append(source + previousTokenEndPosition,
+                            tk.begin() - previousTokenEndPosition);
+            }
+
+            QMultiMap<unsigned, QByteArray>::const_iterator it;
+
+            it = _insertBefore.constFind(i);
+            for (; it != _insertBefore.constEnd() && it.key() == i; ++it) {
+                out->append(it.value());
+            }
+
+            if (! _removed.contains(i))
+                out->append(source + tk.begin(), tk.length);
+
+            it = _insertAfter.constFind(i);
+            for (; it != _insertAfter.constEnd() && it.key() == i; ++it) {
+                out->append(it.value());
+            }
+
+            previousTokenEndPosition = tk.end();
+        }
+    }
+};
+
+class SimpleRefactor: protected ASTVisitor, Rewrite {
+public:
+    SimpleRefactor(Control *control)
+        : ASTVisitor(control)
+    { }
+
+    void operator()(const TranslationUnit *unit,
+                    const QByteArray &source,
+                    QByteArray *out)
+    {
+        accept(unit->ast());
+        rewrite(unit, source, out);
+    }
+
+protected:
+    virtual bool visit(AccessDeclarationAST *ast)
+    {
+        if (tokenKind(ast->access_specifier_token) == T_PRIVATE) {
+            // change visibility from `private' to `public'.
+            remove(ast->access_specifier_token);
+            insertTextAfter(ast->access_specifier_token, "public /* #REF# private->public */");
+        }
+        return true;
+    }
+
+    virtual bool visit(FunctionDefinitionAST *ast)
+    {
+        bool isInline = false;
+        for (SpecifierAST *spec = ast->decl_specifier_seq; spec; spec = spec->next) {
+            if (SimpleSpecifierAST *simpleSpec = spec->asSimpleSpecifier()) {
+                if (tokenKind(simpleSpec->specifier_token) == T_INLINE) {
+                    isInline = true;
+                    break;
+                }
+            }
+        }
+
+        // force the `inline' specifier.
+        if (! isInline)
+            insertTextBefore(ast->firstToken(), "inline /* #REF# made inline */ ");
+
+        return true;
+    }
+
+    virtual bool visit(ClassSpecifierAST *ast)
+    {
+        // export/import the class using the macro MY_EXPORT.
+        if (ast->name)
+            insertTextBefore(ast->name->firstToken(), "MY_EXPORT ");
+
+        // add QObject to the base clause.
+        if (ast->colon_token)
+            insertTextAfter(ast->colon_token, " public QObject,");
+        else if (ast->lbrace_token)
+            insertTextBefore(ast->lbrace_token, ": public QObject ");
+
+        // mark the class as Q_OBJECT.
+        if (ast->lbrace_token)
+            insertTextAfter(ast->lbrace_token, " Q_OBJECT\n");
+
+        return true;
+    }
+};
+
 int main(int, char *[])
 {
-    Control control;
-    StringLiteral *fileId = control.findOrInsertFileName("<stdin>");
-
     QFile in;
     if (! in.open(stdin, QFile::ReadOnly))
         return EXIT_FAILURE;
 
     const QByteArray source = in.readAll();
 
+    Control control;
+    StringLiteral *fileId = control.findOrInsertFileName("<stdin>");
     TranslationUnit unit(&control, fileId);
     unit.setSource(source.constData(), source.size());
     unit.parse();
-    if (unit.ast()) {
-        TranslationUnitAST *ast = unit.ast()->asTranslationUnit();
-        Q_ASSERT(ast != 0);
+    if (! unit.ast())
+        return EXIT_FAILURE;
 
-        Scope globalScope;
-        Semantic sem(&control);
-        for (DeclarationAST *decl = ast->declarations; decl; decl = decl->next) {
-            sem.check(decl, &globalScope);
-        }
+    TranslationUnitAST *ast = unit.ast()->asTranslationUnit();
+    Q_ASSERT(ast != 0);
+
+    Scope globalScope;
+    Semantic sem(&control);
+    for (DeclarationAST *decl = ast->declarations; decl; decl = decl->next) {
+        sem.check(decl, &globalScope);
     }
 
+    // test the rewriter
+    QByteArray out;
+    SimpleRefactor refactor(&control);
+    refactor(&unit, source, &out);
+    printf("%s\n", out.constData());
+
     return EXIT_SUCCESS;
 }