if/elif/else statements using ASTBlocks

Author: dpogue

ASTNode.cpp

@@ -67,7 +67,7 @@ void ASTBlock::removeFirst()
 const char* ASTBlock::type_str() const
 {
     static const char* s_type_strings[] = {
-        "", "try", "except", "finally", "while", "for"
+        "", "if", "else", "elif", "try", "except", "finally", "while", "for"
     };
-    return s_type_strings[type()];
+    return s_type_strings[blktype()];
 }

ASTNode.h

@@ -336,31 +336,44 @@ class ASTBlock : public ASTNode {
 public:
     typedef std::list<PycRef<ASTNode> > list_t;
 
-    enum Type {
-        BLK_MAIN, BLK_TRY, BLK_EXCEPT, BLK_FINALLY, BLK_WHILE, BLK_FOR
+    enum BlkType {
+        BLK_MAIN, BLK_IF, BLK_ELSE, BLK_ELIF, BLK_TRY, BLK_EXCEPT,
+        BLK_FINALLY, BLK_WHILE, BLK_FOR
     };
 
-    ASTBlock(Type type, unsigned int start = 0, unsigned int end = 0)
-        : ASTNode(NODE_BLOCK), m_type(type), m_start(start), m_end(end)
-          { }
+    ASTBlock(BlkType blktype, unsigned int end = 0)
+        : ASTNode(NODE_BLOCK), m_blktype(blktype), m_end(end) { }
 
-    Type type() const { return m_type; }
-    unsigned int start() const { return m_start; }
+    BlkType blktype() const { return m_blktype; }
     unsigned int end() const { return m_end; }
     const list_t& nodes() const { return m_nodes; }
+    list_t::size_type size() const { return m_nodes.size(); }
     void removeFirst();
     void removeLast();
     void append(PycRef<ASTNode> node) { m_nodes.push_back(node); }
     const char* type_str() const;
 
 private:
-    Type m_type;
-    unsigned int m_start;
+    BlkType m_blktype;
     unsigned int m_end;
     list_t m_nodes;
 };
 
-class ASTTryBlock : public ASTBlock {
+class ASTCondBlock : public ASTBlock {
+public:
+    ASTCondBlock(ASTBlock::BlkType blktype, unsigned int end, PycRef<ASTNode> cond,
+            bool negative = false)
+        : ASTBlock(blktype, end), m_cond(cond), m_negative(negative) { }
+
+    PycRef<ASTNode> cond() const { return m_cond; }
+    bool negative() const { return m_negative; }
+
+private:
+    PycRef<ASTNode> m_cond;
+    bool m_negative;
+};
+
+/*class ASTTryBlock : public ASTBlock {
 public:
     ASTTryBlock(unsigned int start, unsigned int end,
             unsigned int except = 0, unsigned int finally = 0)
@@ -376,6 +389,6 @@ class ASTTryBlock : public ASTBlock {
 private:
     unsigned int m_except;
     unsigned int m_finally;
-};
+};*/
 
 #endif

ASTree.cpp

@@ -240,7 +240,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
             break;
         case Pyc::END_FINALLY:
             {
-                if (curblock->type() == ASTBlock::BLK_FINALLY) {
+                if (curblock->blktype() == ASTBlock::BLK_FINALLY) {
                     blocks.pop();
                     blocks.top()->append(curblock.cast<ASTNode>());
 
@@ -291,40 +291,59 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
             break;
         case Pyc::JUMP_IF_FALSE_A:
             {
+                stack_hist.push(stack);
+                //stack = FastStack(stack_hist.top());
                 PycRef<ASTNode> cond = stack.top();
-                stack.pop();
-                stack.push(new ASTJump(operand, ASTJump::JMP_FALSE, cond));
-                jumps.push(pos + operand);
-                startBlock = true;
+                // Do not pop the condition off the stack!
+
+                PycRef<ASTBlock> ifblk;
+
+                if (curblock->blktype() == ASTBlock::BLK_ELSE && curblock->size() == 0) {
+                    blocks.pop();
+                    ifblk = new ASTCondBlock(ASTBlock::BLK_ELIF, pos+operand, cond, false);
+                } else {
+                    ifblk = new ASTCondBlock(ASTBlock::BLK_IF, pos+operand, cond, false);
+                }
+                blocks.push(ifblk.cast<ASTBlock>());
+                curblock = blocks.top();
             }
             break;
         case Pyc::JUMP_IF_TRUE_A:
             {
+                stack_hist.push(stack);
+                //stack = FastStack(stack_hist.top());
                 PycRef<ASTNode> cond = stack.top();
-                stack.pop();
-                stack.push(new ASTJump(operand, ASTJump::JMP_TRUE, cond));
-                jumps.push(pos + operand);
-                startBlock = true;
+                // Do not pop the condition off the stack!
+
+                PycRef<ASTBlock> ifblk;
+
+                if (curblock->blktype() == ASTBlock::BLK_ELSE && curblock->size() == 0) {
+                    blocks.pop();
+                    ifblk = new ASTCondBlock(ASTBlock::BLK_ELIF, pos+operand, cond, true);
+                } else {
+                    ifblk = new ASTCondBlock(ASTBlock::BLK_IF, pos+operand, cond, true);
+                }
+                blocks.push(ifblk.cast<ASTBlock>());
+                curblock = blocks.top();
             }
             break;
         case Pyc::JUMP_FORWARD_A:
             {
-                stack.push(new ASTJump(operand, ASTJump::JUMP, NULL));
-                /*std::stack<int> tmp;
-                int top = jumps.top();
-                while (top < pos + operand) {
-                    tmp.push(jumps.top());
-                    jumps.pop();
-                    if (!jumps.size())
-                        break;
-                    top = jumps.top();
-                }
-                jumps.push(pos + operand);
-                while (tmp.size()) {
-                    jumps.push(tmp.top());
-                    tmp.pop();
+                //stack = FastStack(stack_hist.top());
+                stack_hist.pop();
+
+                PycRef<ASTBlock> prev = curblock;
+                blocks.pop();
+
+                blocks.top()->append(prev.cast<ASTNode>());
+
+                if (operand > 1 && (prev->blktype() == ASTBlock::BLK_IF
+                            || prev->blktype() == ASTBlock::BLK_ELIF))
+                {
+                    PycRef<ASTBlock> next = new ASTBlock(ASTBlock::BLK_ELSE, pos+operand);
+                    blocks.push(next.cast<ASTBlock>());
                 }
-                startBlock = true;*/
+                curblock = blocks.top();
             }
             break;
         case Pyc::LOAD_ATTR_A:
@@ -374,7 +393,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                 stack.push(new ASTFunction(code, defArgs));
             }
             break;
-        case Pyc::POP_BLOCK:
+        /*case Pyc::POP_BLOCK:
             {
                 PycRef<ASTBlock> tmp = curblock;
 
@@ -392,21 +411,15 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                     //todo
                 }
             }
-            break;
+            break;*/
         case Pyc::POP_TOP:
             {
                 PycRef<ASTNode> value = stack.top();
                 stack.pop();
-                if (value->type() == ASTNode::NODE_CALL || value->type() == ASTNode::NODE_JUMP) {
+                if (value->type() == ASTNode::NODE_CALL) {
                     lines.top().push_back(value);
                     curblock->append(value);
                 }
-
-                /*if (startBlock) {
-                    ASTNodeList::list_t blk;
-                    lines.push(blk);
-                    startBlock = false;
-                }*/
             }
             break;
         case Pyc::PRINT_ITEM:
@@ -442,7 +455,7 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
         case Pyc::SET_LINENO_A:
             // Ignore
             break;
-        case Pyc::SETUP_EXCEPT_A:
+        /*case Pyc::SETUP_EXCEPT_A:
             {
                 if (curblock->type() == ASTBlock::BLK_TRY) {
                     PycRef<ASTTryBlock> tryblk = curblock.cast<ASTTryBlock>();
@@ -459,14 +472,14 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
                     curblock = blocks.top();
                 }
             }
-            break;
-        case Pyc::SETUP_FINALLY_A:
+            break;*/
+        /*case Pyc::SETUP_FINALLY_A:
             {
                 PycRef<ASTBlock> tryblk = new ASTTryBlock(pos, pos+operand, 0, pos+operand);
                 blocks.push(tryblk.cast<ASTBlock>());
                 curblock = blocks.top();
             }
-            break;
+            break;*/
         case Pyc::STORE_ATTR_A:
             {
                 PycRef<ASTNode> name = stack.top();
@@ -566,16 +579,15 @@ PycRef<ASTNode> BuildFromCode(PycRef<PycCode> code, PycModule* mod)
             return new ASTNodeList(defblock->nodes());
         }
 
-        /*while (jumps.size() && jumps.top() == pos)
-        {
-            ASTNodeList::list_t block = lines.top();
-            if (lines.size() > 1)
-            {
-                lines.pop();
-            }
-            lines.top().push_back(new ASTNodeList(block));
-            jumps.pop();
-        }*/
+        if (curblock->end() && curblock->end() < pos) {
+            //stack = FastStack(stack_hist.top());
+            stack_hist.pop();
+
+            PycRef<ASTBlock> prev = curblock;
+            blocks.pop();
+            curblock = blocks.top();
+            curblock->append(prev.cast<ASTNode>());
+        }
     }
 
     cleanBuild = true;
@@ -766,15 +778,29 @@ void print_src(PycRef<ASTNode> node, PycModule* mod)
         break;
     case ASTNode::NODE_BLOCK:
         {
-            printf("%s:\n", node.cast<ASTBlock>()->type_str());
+            printf("%s", node.cast<ASTBlock>()->type_str());
+            if (node.cast<ASTBlock>()->blktype() == ASTBlock::BLK_IF
+                    || node.cast<ASTBlock>()->blktype() == ASTBlock::BLK_ELIF)
+            {
+                if (node.cast<ASTCondBlock>()->negative())
+                    printf(" not ");
+                else
+                    printf(" ");
+
+                print_src(node.cast<ASTCondBlock>()->cond(), mod);
+            }
+            printf(":\n");
+
             cur_indent++;
             ASTBlock::list_t lines = node.cast<ASTBlock>()->nodes();
-            for (ASTBlock::list_t::const_iterator ln = lines.begin(); ln != lines.end(); ++ln) {
+            for (ASTBlock::list_t::const_iterator ln = lines.begin(); ln != lines.end();) {
                 if ((*ln).cast<ASTNode>()->type() != ASTNode::NODE_NODELIST) {
                     start_line(cur_indent);
                 }
                 print_src(*ln, mod);
-                end_line();
+                if (++ln != lines.end()) {
+                    end_line();
+                }
             }
             cur_indent--;
         }

object.h

@@ -6,21 +6,27 @@ class PycRef {
 public:
     PycRef() : m_obj(0) { }
     PycRef(_Obj* obj) : m_obj(obj) { m_obj->addRef(); }
-    PycRef(const PycRef<_Obj>& obj) : m_obj(obj.m_obj) { m_obj->addRef(); }
-    ~PycRef<_Obj>() { m_obj->delRef(); }
+    PycRef(const PycRef<_Obj>& obj) : m_obj(obj.m_obj) {
+        if (m_obj != (_Obj*)0) m_obj->addRef();
+    }
+    ~PycRef<_Obj>() { if (m_obj != (_Obj*)0) m_obj->delRef(); }
 
     PycRef<_Obj>& operator=(_Obj* obj)
     {
-        obj->addRef();
-        m_obj->delRef();
+        if (obj != (_Obj*)0)
+            obj->addRef();
+        if (m_obj != (_Obj*)0)
+            m_obj->delRef();
         m_obj = obj;
         return *this;
     }
 
     PycRef<_Obj>& operator=(const PycRef<_Obj>& obj)
     {
-        obj.m_obj->addRef();
-        m_obj->delRef();
+        if (obj.m_obj != (_Obj*)0)
+            obj.m_obj->addRef();
+        if (m_obj != (_Obj*)0)
+            m_obj->delRef();
         m_obj = obj.m_obj;
         return *this;
     }