Fixing bug with final branch in AST