Fix bug in dce

src/dead_code_elimination.cpp

@@ -3,9 +3,29 @@
 #include <migraph/instruction.hpp>
 #include <migraph/iterator_for.hpp>
 #include <migraph/functional.hpp>
+#include <migraph/ranges.hpp>
 
 namespace migraph {
 
+template<class Range, class Iterator>
+std::ptrdiff_t bidistance(const Range& r, Iterator start, Iterator last)
+{
+    auto start_forward = start;
+    auto start_backwards = start;
+    std::size_t n = 0;
+    while(start_forward != last and start_backwards != last)
+    {
+        n++;
+        if(start_forward != r.end()) start_forward++;
+        if(start_backwards != r.begin()) start_backwards--;
+    }
+    if(start_forward == last)
+        return n;
+    else
+        return -n;
+
+}
+
 void dead_code_elimination::apply(program& p) const
 {
     auto last = std::prev(p.end());
@@ -16,18 +36,21 @@ void dead_code_elimination::apply(program& p) const
         if(ins == p.begin())
             continue;
         const auto i = std::prev(ins);
-        // Skip instruction with empty shape as output unless its a builtin
-        if(i->get_shape().elements() == 0 and not(i->name().front() == '@'))
-            continue;
         // Skip the last instruction
         if(i == last)
             break;
+        // Skip instruction with empty shape as output unless its a builtin
+        if(i->get_shape().elements() == 0 and not(i->name().front() == '@'))
+            continue;
+        assert(bidistance(p, i, last) > 0);
         fix([&](auto self, auto leaf) {
             assert(p.has_instruction(leaf));
             if(leaf->outputs().empty())
             {
                 auto args = leaf->inputs();
                 leaf->clear_arguments();
+                assert(bidistance(p, last, leaf) < 0);
+                assert(leaf != ins);
                 p.move_instruction(leaf, p.end());
                 for(auto arg : args)
                     self(arg);

test/dead_code_elimination_test.cpp

@@ -43,6 +43,23 @@ void simple_test_nop()
     EXPECT(result != migraph::literal{4});
 }
 
+void simple_test_nop2()
+{
+    migraph::program p;
+
+    auto one = p.add_literal(1);
+    auto two = p.add_literal(2);
+    p.add_instruction(nop{});
+    p.add_instruction(sum_op{}, one, two);
+    p.add_instruction(nop{});
+    auto count = std::distance(p.begin(), p.end());
+    p.compile(dce_target{});
+    EXPECT(std::distance(p.begin(), p.end()) == 2);
+    auto result = p.eval({});
+    EXPECT(result == migraph::literal{});
+    EXPECT(result != migraph::literal{4});
+}
+
 void duplicate_test1()
 {
     migraph::program p;
@@ -99,6 +116,7 @@ int main()
 {
     simple_test();
     simple_test_nop();
+    simple_test_nop2();
     duplicate_test1();
     duplicate_test2();
     depth_test();