Commit b6a4c079 authored by Nguyen, Thien Minh's avatar Nguyen, Thien Minh
Browse files

Support Python conditional statemements 



Signed-off-by: default avatarThien Nguyen <nguyentm@ornl.gov>
parent f415a46a

handlers/token_collector/pyxasm/pyxasm_token_collector.cpp

+40 −4
Original line number Diff line number Diff line
@@ -82,7 +82,9 @@ void PyXasmTokenCollector::collect(clang::Preprocessor &PP, 
    }



    // If statement:

    if (Toks[i].is(clang::tok::TokenKind::kw_if)) {

    // Note: Python has an "elif" token, which doesn't have a C++ equiv.

    if (Toks[i].is(clang::tok::TokenKind::kw_if) ||

        PP.getSpelling(Toks[i]) == "elif") {

      line += " ";

      i += 1;

      line += PP.getSpelling(Toks[i]);
@@ -118,16 +120,50 @@ void PyXasmTokenCollector::collect(clang::Preprocessor &PP, 
      nb_closing_scopes++;

    }



    std::string lineText = line.first;

    // Enter a new for scope block (for/if/etc.) -> push to the stack

    // Note: we rewrite Python if .. elif .. else as follows:

    // Python:

    // if (cond1):

    //   code1

    // elif (cond2):

    //   code2

    // else:

    //   code3

    // ===============

    // C++:

    // if (cond1) {

    //   code1

    // }

    // else if (cond2) {

    //   code2

    // }

    // else {

    //   code3

    // }



    if (line.first.find("for ") != std::string::npos ||

        line.first.find("if ") != std::string::npos) {

        // Starts with 'if'

        line.first.rfind("if ", 0) == 0) {

      scope_block_indent.push(line.second);

    } else if (line.first == "else:") {

      ss << "else {\n";

      scope_block_indent.push(line.second);

    }



    // Starts with 'elif'

    else if (line.first.rfind("elif ", 0) == 0) {

      // Rewrite it to

      // else if () { }

      ss << "else ";

      scope_block_indent.push(line.second);

      // Remove the first two characters ("el")

      // hence this line will be parsed as an idependent C++ if block:

      lineText.erase(0, 2);

    }

    // is_in_for_loop = line.first.find("for ") != std::string::npos &&

    // line.second >= previous_col;



    ANTLRInputStream input(line.first);

    ANTLRInputStream input(lineText);

    pyxasmLexer lexer(&input);

    CommonTokenStream tokens(&lexer);

    pyxasmParser parser(&tokens);

python/tests/test_kernel_jit.py

+31 −0
Original line number Diff line number Diff line
@@ -326,6 +326,37 @@ class TestSimpleKernelJIT(unittest.TestCase): 
        self.assertAlmostEqual((float)(comp.getInstruction(2).getParameter(0)), -1.234)

        self.assertEqual(comp.getInstruction(3).name(), "CNOT") 



    # Test conditional if..elif..else rewrite

    def test_if_clause(self):

        @qjit

        def test_if_stmt(q : qreg, flag: int):

            H(q[0])

            if flag == 0:

                X(q[0])

            elif flag == 1:

                Y(q[0])

            elif flag == 2:

                Z(q[0])

            else:

                T(q[0])

        

        q = qalloc(2)



        # Examine the circuit QASM with various values of flag

        comp0 = test_if_stmt.extract_composite(q, 0)

        comp1 = test_if_stmt.extract_composite(q, 1)

        comp2 = test_if_stmt.extract_composite(q, 2)

        comp3 = test_if_stmt.extract_composite(q, 3)



        self.assertEqual(comp0.nInstructions(), 2)   

        self.assertEqual(comp1.nInstructions(), 2)   

        self.assertEqual(comp2.nInstructions(), 2)   

        self.assertEqual(comp3.nInstructions(), 2)   



        self.assertEqual(comp0.getInstruction(1).name(), "X") 

        self.assertEqual(comp1.getInstruction(1).name(), "Y") 

        self.assertEqual(comp2.getInstruction(1).name(), "Z") 

        self.assertEqual(comp3.getInstruction(1).name(), "T") 



if __name__ == '__main__':

  unittest.main()
 
 No newline at end of file