// See also: file:///usr/share/doc/python/html/reference/grammar.html?highlight=grammar
// See also: https://github.com/antlr/grammars-v4/blob/master/python3/Python3.g4
// See also: file:///usr/share/doc/python/html/reference/compound_stmts.html#function-definitions
// See also: https://greentreesnakes.readthedocs.io/en/latest/nodes.html#keyword

use crate::{
    ast,
    lexer::{LexicalError, LexicalErrorType},
    function::{ArgumentList, parse_args, parse_params, validate_arguments},
    context::set_context,
    string::parse_strings,
    token::{self, StringKind},
};
use num_bigint::BigInt;

grammar;

// This is a hack to reduce the amount of lalrpop tables generated:
// For each public entry point, a full parse table is generated.
// By having only a single pub function, we reduce this to one.
pub Top: ast::Mod = {
    StartModule <body:Program> => ast::Mod::Module { body, type_ignores: vec![] },
    StartInteractive <body:Program> => ast::Mod::Interactive { body },
    StartExpression <body:TestList> ("\n")* => ast::Mod::Expression { body: Box::new(body) },
};

Program: ast::Suite = {
    <lines:FileLine*> => {
        lines.into_iter().flatten().collect()
    },
};

// A file line either has a declaration, or an empty newline:
FileLine: ast::Suite = {
    Statement,
    "\n" => vec![],
};

Suite: ast::Suite = {
    SimpleStatement,
    "\n" Indent <s:Statement+> Dedent => s.into_iter().flatten().collect(),
};

Statement: ast::Suite = {
    SimpleStatement,
    <s:CompoundStatement> => vec![s],
};

SimpleStatement: ast::Suite = {
    <s1:SmallStatement> <s2:(";" SmallStatement)*> ";"? "\n" => {
        let mut statements = vec![s1];
        statements.extend(s2.into_iter().map(|e| e.1));
        statements
    }
};

SmallStatement: ast::Stmt = {
    ExpressionStatement,
    PassStatement,
    DelStatement,
    FlowStatement,
    ImportStatement,
    GlobalStatement,
    NonlocalStatement,
    AssertStatement,
};

PassStatement: ast::Stmt = {
    <location:@L> "pass" <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Pass,
        )
    },
};

DelStatement: ast::Stmt = {
    <location:@L> "del" <targets:ExpressionList2> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Delete { targets: targets.into_iter().map(|expr| set_context(expr, ast::ExprContext::Del)).collect() },
        )
    },
};

ExpressionStatement: ast::Stmt = {
    <location:@L> <expression:TestOrStarExprList> <suffix:AssignSuffix*> <end_location:@R> => {
        // Just an expression, no assignment:
        if suffix.is_empty() {
            ast::Stmt::new(
                location,
                end_location,
                ast::StmtKind::Expr { value: Box::new(expression) }
            )
        } else {
            let mut targets = vec![set_context(expression, ast::ExprContext::Store)];
            let mut values = suffix;

            while values.len() > 1 {
                targets.push(set_context(values.remove(0), ast::ExprContext::Store));
            }

            let value = Box::new(values.into_iter().next().unwrap());

            ast::Stmt::new(
                location,
                end_location,
                ast::StmtKind::Assign { targets, value, type_comment: None },
            )
        }
    },
    <location:@L> <target:TestOrStarExprList> <op:AugAssign> <rhs:TestListOrYieldExpr> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::AugAssign {
                target: Box::new(set_context(target, ast::ExprContext::Store)),
                op,
                value: Box::new(rhs)
            },
        )
    },
    <location:@L> <target:Test<"all">> ":" <annotation:Test<"all">> <rhs:AssignSuffix?> <end_location:@R> => {
        let simple = matches!(target.node, ast::ExprKind::Name { .. });
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::AnnAssign {
                target: Box::new(set_context(target, ast::ExprContext::Store)),
                annotation: Box::new(annotation),
                value: rhs.map(Box::new),
                simple: if simple { 1 } else { 0 },
            },
        )
    },
};

AssignSuffix: ast::Expr = {
    "=" <e:TestListOrYieldExpr> => e
};

TestListOrYieldExpr: ast::Expr = {
    TestList,
    YieldExpr
}

#[inline]
TestOrStarExprList: ast::Expr = {
    // as far as I can tell, these were the same
    TestList
};

TestOrStarExpr: ast::Expr = {
    Test<"all">,
    StarExpr,
};

NamedOrStarExpr: ast::Expr = {
    NamedExpression,
    StarExpr,
};

TestOrStarNamedExpr: ast::Expr = {
    NamedExpressionTest,
    StarExpr,
};

AugAssign: ast::Operator = {
    "+=" => ast::Operator::Add,
    "-=" => ast::Operator::Sub,
    "*=" => ast::Operator::Mult,
    "@=" => ast::Operator::MatMult,
    "/=" => ast::Operator::Div,
    "%=" => ast::Operator::Mod,
    "&=" => ast::Operator::BitAnd,
    "|=" => ast::Operator::BitOr,
    "^=" => ast::Operator::BitXor,
    "<<=" => ast::Operator::LShift,
    ">>=" => ast::Operator::RShift,
    "**=" => ast::Operator::Pow,
    "//=" => ast::Operator::FloorDiv,
};

FlowStatement: ast::Stmt = {
    <location:@L> "break" <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Break,
        )
    },
    <location:@L> "continue" <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Continue,
        )
    },
    <location:@L> "return" <value:TestList?> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Return { value: value.map(Box::new) },
        )
    },
    <location:@L> <expression:YieldExpr> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Expr { value: Box::new(expression) },
        )
    },
    RaiseStatement,
};

RaiseStatement: ast::Stmt = {
    <location:@L> "raise" <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Raise { exc: None, cause: None },
        )
    },
    <location:@L> "raise" <t:Test<"all">> <c:("from" Test<"all">)?> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Raise { exc: Some(Box::new(t)), cause: c.map(|x| Box::new(x.1)) },
        )
    },
};

ImportStatement: ast::Stmt = {
    <location:@L> "import" <names: OneOrMore<ImportAsAlias<DottedName>>> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Import { names },
        )
    },
    <location:@L> "from" <source:ImportFromLocation> "import" <names: ImportAsNames> <end_location:@R> => {
        let (level, module) = source;
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::ImportFrom {
                level,
                module,
                names
            },
        )
    },
};

ImportFromLocation: (Option<usize>, Option<String>) = {
    <dots: ImportDots*> <name:DottedName> => {
        (Some(dots.iter().sum()), Some(name))
    },
    <dots: ImportDots+> => {
        (Some(dots.iter().sum()), None)
    },
};

ImportDots: usize = {
    "..." => 3,
    "." => 1,
};

ImportAsNames: Vec<ast::Alias> = {
    <location:@L> <i:OneOrMore<ImportAsAlias<Identifier>>> <end_location:@R> => i,
    <location:@L> "(" <i:OneOrMore<ImportAsAlias<Identifier>>> ","? ")" <end_location:@R> => i,
    <location:@L> "*" <end_location:@R> => {
        // Star import all
        vec![ast::Alias::new(location, end_location, ast::AliasData { name: "*".to_string(), asname: None })]
    },
};


#[inline]
ImportAsAlias<I>: ast::Alias = {
    <location:@L> <name:I> <a: ("as" Identifier)?> <end_location:@R> => ast::Alias::new(location, end_location, ast::AliasData { name, asname: a.map(|a| a.1) }),
}

// A name like abc or abc.def.ghi
DottedName: String = {
    <n:name> => n,
    <n:name> <n2: ("." Identifier)+>  => {
        let mut r = n.to_string();
        for x in n2 {
            r.push_str(".");
            r.push_str(&x.1);
        }
        r
    },
};

GlobalStatement: ast::Stmt = {
    <location:@L> "global" <names:OneOrMore<Identifier>> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Global { names }
        )
    },
};

NonlocalStatement: ast::Stmt = {
    <location:@L> "nonlocal" <names:OneOrMore<Identifier>> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Nonlocal { names }
        )
    },
};

AssertStatement: ast::Stmt = {
    <location:@L> "assert" <test:Test<"all">> <msg: ("," Test<"all">)?> <end_location:@R> => {
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Assert {
                test: Box::new(test),
                msg: msg.map(|e| Box::new(e.1))
            }
        )
    },
};

CompoundStatement: ast::Stmt = {
    MatchStatement,
    IfStatement,
    WhileStatement,
    ForStatement,
    TryStatement,
    WithStatement,
    FuncDef,
    ClassDef,
};

MatchStatement: ast::Stmt = {
    <location:@L> "match" <subject:TestOrStarNamedExpr> ":" "\n" Indent <cases:MatchCase+> Dedent => {
       let end_location = cases
            .last()
            .unwrap()
            .body
            .last()
            .unwrap()
            .end();
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Match {
                subject: Box::new(subject),
                cases
            }
        )
    },
    <location:@L> "match" <subject:TestOrStarNamedExpr> "," ":" "\n" Indent <cases:MatchCase+> Dedent => {
       let end_location = cases
            .last()
            .unwrap()
            .body
            .last()
            .unwrap()
            .end();
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Match {
                subject: Box::new(subject),
                cases
            }
        )
    },
    <location:@L> "match" <subject:TestOrStarNamedExpr> "," <subjects:OneOrMore<TestOrStarNamedExpr>> ","? ":" "\n" Indent <cases:MatchCase+> Dedent => {
        let end_location = cases
            .last()
            .unwrap()
            .body
            .last()
            .unwrap()
            .end();
        let mut subjects = subjects;
        subjects.insert(0, subject);
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Match {
                subject: Box::new(ast::Expr::new(
                    location,
                    end_location,
                    ast::ExprKind::Tuple {
                        elts: subjects,
                        ctx: ast::ExprContext::Load,
                    },
                )),
                cases
            }
        )
    }
}

MatchCase: ast::MatchCase = {
    "case" <pattern:Patterns> <guard:(Guard)?> ":" <body:Suite> => {
        ast::MatchCase {
            pattern,
            guard: guard.map(Box::new),
            body
        }
    },
}

Guard: ast::Expr = {
    "if" <guard:NamedExpressionTest> => {
        guard
    }
}

Patterns: ast::Pattern = {
    <location:@L> <pattern:Pattern> "," <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        ast::PatternKind::MatchSequence {
            patterns: vec![pattern]
        },
    ),
    <location:@L> <pattern:Pattern> "," <patterns:OneOrMore<Pattern>> ","? <end_location:@R> => {
        let mut patterns = patterns;
        patterns.insert(0, pattern);
        ast::Pattern::new(
            location,
            end_location,
            ast::PatternKind::MatchSequence {
                patterns
            },
        )
    },
    <pattern:Pattern> => pattern
}

Pattern: ast::Pattern = {
    <pattern:AsPattern> => pattern,
    <pattern:OrPattern> => pattern,
}

AsPattern: ast::Pattern = {
    <location:@L> <pattern:OrPattern> "as" <name:Identifier> <end_location:@R> =>? {
        if name == "_" {
            Err(LexicalError {
                error: LexicalErrorType::OtherError("cannot use '_' as a target".to_string()),
                location,
            })?
        } else {
            Ok(ast::Pattern::new(
                location,
                end_location,
                ast::PatternKind::MatchAs {
                   pattern: Some(Box::new(pattern)),
                   name: Some(name),
               },
            ))
        }
    },
}

OrPattern: ast::Pattern = {
    <pattern:ClosedPattern> => pattern,
    <location:@L> <pattern:ClosedPattern> <patterns:("|" <ClosedPattern>)+> <end_location:@R> => {
        let mut patterns = patterns;
        patterns.insert(0, pattern);
        ast::Pattern::new(
            location,
            end_location,
            ast::PatternKind::MatchOr { patterns }
        )
    }
}

ClosedPattern: ast::Pattern = {
    <location:@L> <node:LiteralPattern> <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        node,
    ),
    <location:@L> <node:CapturePattern> <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        node,
    ),
    <location:@L> <node:StarPattern> <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        node,
    ),
    <location:@L> <node:ValuePattern> <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        node,
    ),
    <location:@L> <node:SequencePattern> <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        node,
    ),
    <location:@L> <node:MappingPattern> <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        node,
    ),
    <location:@L> <node:ClassPattern> <end_location:@R> => ast::Pattern::new(
        location,
        end_location,
        node,
    ),
}

SequencePattern: ast::PatternKind = {
    // A single-item tuple is a special case: it's a group pattern, _not_ a sequence pattern.
    <location:@L> "(" <pattern:Pattern> ")" <end_location:@R> => pattern.node,
    <location:@L> "(" ")" <end_location:@R> => ast::PatternKind::MatchSequence {
        patterns: vec![],
    },
    <location:@L> "(" <pattern:Pattern> "," <patterns:Comma<Pattern>> ")" <end_location:@R> => {
        let mut patterns = patterns;
        patterns.insert(0, pattern);
        ast::PatternKind::MatchSequence {
            patterns
        }
    },
    <location:@L> "[" <patterns:Comma<Pattern>> "]" <end_location:@R> => ast::PatternKind::MatchSequence {
        patterns
    },
}

StarPattern: ast::PatternKind = {
    <location:@L> "*" <name:Identifier> <end_location:@R> => ast::PatternKind::MatchStar {
        name: if name == "_" { None } else { Some(name) }
    },
}

ConstantAtom: ast::Expr = {
    <location:@L> <value:Constant> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Constant { value, kind: None }
    ),
}

ConstantExpr: ast::Expr = {
    ConstantAtom,
    <location:@L> "-" <operand:ConstantAtom> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::UnaryOp {
            op: ast::Unaryop::USub,
            operand: Box::new(operand)
        }
    ),
}

AddOpExpr: ast::Expr = {
    <location:@L> <left:ConstantExpr> <op:AddOp> <right:ConstantAtom> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp {
            left: Box::new(left),
            op,
            right: Box::new(right),
        }
    ),
}

LiteralPattern: ast::PatternKind = {
    "None" => ast::PatternKind::MatchSingleton {
        value: ast::Constant::None
    },
    "True" => ast::PatternKind::MatchSingleton {
        value: true.into()
    },
    "False" => ast::PatternKind::MatchSingleton {
        value: false.into()
    },
    <location:@L> <value:ConstantExpr> <end_location:@R> => ast::PatternKind::MatchValue {
        value: Box::new(value)
    },
    <location:@L> <value:AddOpExpr> => ast::PatternKind::MatchValue {
        value: Box::new(value)
    },
    <location:@L> <s:(@L string @R)+> =>? Ok(ast::PatternKind::MatchValue {
        value: Box::new(parse_strings(s)?)
    }),
}

CapturePattern: ast::PatternKind = {
    <location:@L> <name:Identifier> <end_location:@R> => ast::PatternKind::MatchAs {
        pattern: None,
        name: if name == "_" { None } else { Some(name) }
    },
}

MatchName: ast::Expr = {
    <location:@L> <name:Identifier> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Name { id: name, ctx: ast::ExprContext::Load },
    ),
}

MatchNameOrAttr: ast::Expr = {
    <location:@L> <name:MatchName> "." <attr:Identifier> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Attribute {
            value: Box::new(name),
            attr,
            ctx: ast::ExprContext::Load,
        },
    ),
    <location:@L> <e:MatchNameOrAttr> "." <attr:Identifier> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Attribute {
            value: Box::new(e),
            attr,
            ctx: ast::ExprContext::Load,
        }
    )
}

ValuePattern: ast::PatternKind = {
   <e:MatchNameOrAttr> => ast::PatternKind::MatchValue {
       value: Box::new(e)
   },
}

MappingKey: ast::Expr = {
    ConstantExpr,
    AddOpExpr,
    MatchNameOrAttr,
    <location:@L> "None" <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Constant {
            value: ast::Constant::None,
            kind: None,
        },
    ),
    <location:@L> "True" <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Constant {
            value: true.into(),
            kind: None,
        },
    ),
    <location:@L> "False" <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Constant {
            value: false.into(),
            kind: None,
        },
    ),
    <location:@L> <s:(@L string @R)+> =>? Ok(parse_strings(s)?),
}

MatchMappingEntry: (ast::Expr, ast::Pattern) = {
    <k:MappingKey> ":" <v:Pattern> => (k, v),
};

MappingPattern: ast::PatternKind = {
    <location:@L> "{" "}" <end_location:@R> => {
        return ast::PatternKind::MatchMapping {
            keys: vec![],
            patterns: vec![],
            rest: None,
        };
    },
    <location:@L> "{" <e:OneOrMore<MatchMappingEntry>> ","? "}" <end_location:@R> => {
        let (keys, patterns) = e
            .into_iter()
            .unzip();
        return ast::PatternKind::MatchMapping {
            keys,
            patterns,
            rest: None,
        };
    },
    <location:@L> "{" "**" <rest:Identifier> ","? "}" <end_location:@R> => {
        return ast::PatternKind::MatchMapping {
            keys: vec![],
            patterns: vec![],
            rest: Some(rest),
        };
    },
    <location:@L> "{" <e:OneOrMore<MatchMappingEntry>> ","  "**" <rest:Identifier> ","? "}" <end_location:@R> => {
        let (keys, patterns) = e
            .into_iter()
            .unzip();
        return ast::PatternKind::MatchMapping {
            keys,
            patterns,
            rest: Some(rest),
        };
    },
}

MatchKeywordEntry: (String, ast::Pattern) = {
    <k:Identifier> "=" <v:Pattern> => (k, v),
};

ClassPattern: ast::PatternKind = {
    <location:@L> <e:MatchName> "(" <patterns: OneOrMore<Pattern>> "," <kwds:OneOrMore<MatchKeywordEntry>> ","? ")" <end_location:@R> => {
        let (kwd_attrs, kwd_patterns) = kwds
            .into_iter()
            .unzip();
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns,
            kwd_attrs,
            kwd_patterns,
        }
    },
    <location:@L> <e:MatchName> "(" <patterns: OneOrMore<Pattern>> ","? ")" <end_location:@R> => {
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns,
            kwd_attrs: vec![],
            kwd_patterns: vec![],
        }
    },
    <location:@L> <e:MatchName> "(" <kwds:OneOrMore<MatchKeywordEntry>> ","? ")" <end_location:@R> => {
         let (kwd_attrs, kwd_patterns) = kwds
            .into_iter()
            .unzip();
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns: vec![],
            kwd_attrs,
            kwd_patterns,
        }
    },
    <location:@L> <e:MatchName> "(" ")" <end_location:@R> => {
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns: vec![],
            kwd_attrs: vec![],
            kwd_patterns: vec![],
        }
    },
    <location:@L> <e:MatchNameOrAttr> "(" <patterns: OneOrMore<Pattern>> "," <kwds:OneOrMore<MatchKeywordEntry>> ","? ")" <end_location:@R> => {
        let (kwd_attrs, kwd_patterns) = kwds
            .into_iter()
            .unzip();
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns,
            kwd_attrs,
            kwd_patterns,
        }
    },
    <location:@L> <e:MatchNameOrAttr> "(" <patterns: OneOrMore<Pattern>> ","? ")" <end_location:@R> => {
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns,
            kwd_attrs: vec![],
            kwd_patterns: vec![],
        }
    },
    <location:@L> <e:MatchNameOrAttr> "(" <kwds:OneOrMore<MatchKeywordEntry>> ","? ")" <end_location:@R> => {
         let (kwd_attrs, kwd_patterns) = kwds
            .into_iter()
            .unzip();
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns: vec![],
            kwd_attrs,
            kwd_patterns,
        }
    },
    <location:@L> <e:MatchNameOrAttr> "(" ")" <end_location:@R> => {
        ast::PatternKind::MatchClass {
            cls: Box::new(e),
            patterns: vec![],
            kwd_attrs: vec![],
            kwd_patterns: vec![],
        }
    },
}

IfStatement: ast::Stmt = {
    <location:@L> "if" <test:NamedExpressionTest> ":" <body:Suite> <s2:(@L "elif" NamedExpressionTest ":" Suite)*> <s3:("else" ":" Suite)?> => {
        // Determine last else:
        let mut last = s3.map(|s| s.2).unwrap_or_default();
        let end_location = last
            .last()
            .or_else(|| s2.last().and_then(|last| last.4.last()))
            .or_else(|| body.last())
            .unwrap()
            .end();
        // handle elif:
        for i in s2.into_iter().rev() {
            let x = ast::Stmt::new(
                i.0,
                end_location,
                ast::StmtKind::If { test: Box::new(i.2), body: i.4, orelse: last },
            );
            last = vec![x];
        }

        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::If { test: Box::new(test), body, orelse: last }
        )
    },
};

WhileStatement: ast::Stmt = {
    <location:@L> "while" <test:NamedExpressionTest> ":" <body:Suite> <s2:("else" ":" Suite)?> => {
        let orelse = s2.map(|s| s.2).unwrap_or_default();
        let end_location = orelse
            .last()
            .or_else(|| body.last())
            .unwrap()
            .end();
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::While {
                test: Box::new(test),
                body,
                orelse
            },
        )
    },
};

ForStatement: ast::Stmt = {
    <location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:TestList> ":" <body:Suite> <s2:("else" ":" Suite)?> => {
        let orelse = s2.map(|s| s.2).unwrap_or_default();
        let end_location = orelse
            .last()
            .or_else(|| body.last())
            .unwrap()
            .end();
        let target = Box::new(set_context(target, ast::ExprContext::Store));
        let iter = Box::new(iter);
        let type_comment = None;
        let node = if is_async.is_some() {
            ast::StmtKind::AsyncFor { target, iter, body, orelse, type_comment }
        } else {
            ast::StmtKind::For { target, iter, body, orelse, type_comment }
        };
        ast::Stmt::new(location, end_location, node)
    },
};

TryStatement: ast::Stmt = {
    <location:@L> "try" ":" <body:Suite> <handlers:ExceptClause+> <else_suite:("else" ":" Suite)?> <finally:("finally" ":" Suite)?> <end_location:@R> => {
        let orelse = else_suite.map(|s| s.2).unwrap_or_default();
        let finalbody = finally.map(|s| s.2).unwrap_or_default();
        let end_location = finalbody
            .last()
            .map(|last| last.end())
            .or_else(|| orelse.last().map(|last| last.end()))
            .or_else(|| handlers.last().map(|last| last.end()))
            .unwrap();
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Try {
                body,
                handlers,
                orelse,
                finalbody,
            },
        )
    },
    <location:@L> "try" ":" <body:Suite> <handlers:ExceptStarClause+> <else_suite:("else" ":" Suite)?> <finally:("finally" ":" Suite)?> <end_location:@R> => {
        let orelse = else_suite.map(|s| s.2).unwrap_or_default();
        let finalbody = finally.map(|s| s.2).unwrap_or_default();
        let end_location = finalbody
            .last()
            .or_else(|| orelse.last())
            .map(|last| last.end())
            .or_else(|| handlers.last().map(|last| last.end()))
            .unwrap();
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::TryStar {
                body,
                handlers,
                orelse,
                finalbody,
            },
        )
    },
    <location:@L> "try" ":" <body:Suite> <finally:("finally" ":" Suite)> => {
        let handlers = vec![];
        let orelse = vec![];
        let finalbody = finally.2;
        let end_location = finalbody.last().unwrap().end();
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::Try {
                body,
                handlers,
                orelse,
                finalbody,
            },
        )
    },
};

ExceptStarClause: ast::Excepthandler = {
    <location:@L> "except" "*" <typ:Test<"all">> ":" <body:Suite> => {
        let end_location = body.last().unwrap().end();
        ast::Excepthandler::new(
            location,
            end_location,
            ast::ExcepthandlerKind::ExceptHandler {
                type_: Some(Box::new(typ)),
                name: None,
                body,
            },
        )
    },
    <location:@L> "except" "*" <x:(Test<"all"> "as" Identifier)> ":" <body:Suite> => {
        let end_location = body.last().unwrap().end();
        ast::Excepthandler::new(
            location,
            end_location,
            ast::ExcepthandlerKind::ExceptHandler {
                type_: Some(Box::new(x.0)),
                name: Some(x.2),
                body,
            },
        )
    },
};


ExceptClause: ast::Excepthandler = {
    <location:@L> "except" <typ:Test<"all">?> ":" <body:Suite> => {
        let end_location = body.last().unwrap().end();
        ast::Excepthandler::new(
            location,
            end_location,
            ast::ExcepthandlerKind::ExceptHandler {
                type_: typ.map(Box::new),
                name: None,
                body,
            },
        )
    },
    <location:@L> "except" <x:(Test<"all"> "as" Identifier)> ":" <body:Suite> => {
        let end_location = body.last().unwrap().end();
        ast::Excepthandler::new(
            location,
            end_location,
            ast::ExcepthandlerKind::ExceptHandler {
                type_: Some(Box::new(x.0)),
                name: Some(x.2),
                body,
            },
        )
    },
};

WithStatement: ast::Stmt = {
    <location:@L> <is_async:"async"?> "with" <items:WithItems> ":" <body:Suite> => {
        let end_location = body.last().unwrap().end();
        let type_comment = None;
        let node = if is_async.is_some() {
            ast::StmtKind::AsyncWith { items, body, type_comment }
        } else {
            ast::StmtKind::With { items, body, type_comment }
        };
        ast::Stmt::new(location, end_location, node)
    },
};

WithItems: Vec<ast::Withitem> = {
    "(" <WithItemsNoAs> ","? ")",
    "(" <left:(<WithItemsNoAs> ",")?> <mid:WithItem<"as">> <right:("," <WithItem<"all">>)*> ","? ")" => {
        left.into_iter().flatten().chain([mid]).chain(right).collect()
    },
    <WithItem<"no-withitems">> => vec![<>],
    <item:WithItem<"all">> <items:("," <WithItem<"all">>)+> => {
        [item].into_iter().chain(items).collect()
    }
};

#[inline]
WithItemsNoAs: Vec<ast::Withitem> = {
    <OneOrMore<Test<"all">>> => {
        <>.into_iter().map(|context_expr| ast::Withitem { context_expr, optional_vars: None }).collect()
    },
}

WithItem<Goal>: ast::Withitem = {
    <Test<Goal>> if Goal != "as" => ast::Withitem { context_expr: <>, optional_vars: None },
    <context_expr:Test<"all">> "as" <vars:Expression<"all">> => {
        let optional_vars = Some(Box::new(set_context(vars, ast::ExprContext::Store)));
        ast::Withitem { context_expr, optional_vars }
    },
};

FuncDef: ast::Stmt = {
    <decorator_list:Decorator*> <location:@L> <is_async:"async"?> "def" <name:Identifier> <args:Parameters> <r:("->" Test<"all">)?> ":" <body:Suite>  => {
        let args = Box::new(args);
        let returns = r.map(|x| Box::new(x.1));
        let end_location = body.last().unwrap().end();
        let type_comment = None;
        let node = if is_async.is_some() {
            ast::StmtKind::AsyncFunctionDef { name, args, body, decorator_list, returns, type_comment }
        } else {
            ast::StmtKind::FunctionDef { name, args, body, decorator_list, returns, type_comment }
        };
        ast::Stmt::new(location, end_location, node)
    },
};

Parameters: ast::Arguments = {
    "(" <a: (ParameterList<TypedParameter, StarTypedParameter>)?> ")" =>? {
        let args = validate_arguments(
            a.unwrap_or_else(|| ast::Arguments {
                posonlyargs: vec![],
                args: vec![],
                vararg: None,
                kwonlyargs: vec![],
                kw_defaults: vec![],
                kwarg: None,
                defaults: vec![]
            })
        )?;

        Ok(args)
    }
};

// Note that this is a macro which is used once for function defs, and
// once for lambda defs.
ParameterList<ArgType, StarArgType>: ast::Arguments = {
    <param1:ParameterDefs<ArgType>> <args2:("," ParameterListStarArgs<ArgType, StarArgType>)?> ","? =>? {
        let (posonlyargs, args, defaults) = parse_params(param1)?;

        // Now gather rest of parameters:
        let (vararg, kwonlyargs, kw_defaults, kwarg) = args2.map_or((None, vec![], vec![], None), |x| x.1);

        Ok(ast::Arguments {
            posonlyargs,
            args,
            kwonlyargs,
            vararg,
            kwarg,
            defaults,
            kw_defaults,
        })
    },
    <param1:ParameterDefs<ArgType>> <kw:("," KwargParameter<ArgType>)> ","? =>? {
        let (posonlyargs, args, defaults) = parse_params(param1)?;

        // Now gather rest of parameters:
        let vararg = None;
        let kwonlyargs = vec![];
        let kw_defaults = vec![];
        let kwarg = kw.1;

        Ok(ast::Arguments {
            posonlyargs,
            args,
            kwonlyargs,
            vararg,
            kwarg,
            defaults,
            kw_defaults,
        })
    },
    <params:ParameterListStarArgs<ArgType, StarArgType>> ","? => {
        let (vararg, kwonlyargs, kw_defaults, kwarg) = params;
        ast::Arguments {
            posonlyargs: vec![],
            args: vec![],
            kwonlyargs,
            vararg,
            kwarg,
            defaults: vec![],
            kw_defaults,
        }
    },
    <kwarg:KwargParameter<ArgType>> ","? => {
        ast::Arguments {
            posonlyargs: vec![],
            args: vec![],
            kwonlyargs: vec![],
            vararg: None,
            kwarg,
            defaults: vec![],
            kw_defaults: vec![],
        }
    },
};

// Use inline here to make sure the "," is not creating an ambiguity.
#[inline]
ParameterDefs<ArgType>: (Vec<(ast::Arg, Option<ast::Expr>)>, Vec<(ast::Arg, Option<ast::Expr>)>) = {
    <args:OneOrMore<ParameterDef<ArgType>>> => {
        (vec![], args)
    },
    <pos_args:OneOrMore<ParameterDef<ArgType>>> "," "/" <args:("," ParameterDef<ArgType>)*>  => {
        (pos_args, args.into_iter().map(|e| e.1).collect())
    },
};

ParameterDef<ArgType>: (ast::Arg, Option<ast::Expr>) = {
    <i:ArgType> => (i, None),
    <i:ArgType> "=" <e:Test<"all">> => (i, Some(e)),
};

UntypedParameter: ast::Arg = {
    <location:@L> <arg:Identifier> <end_location:@R> => ast::Arg::new(
        location,
        end_location,
        ast::ArgData { arg, annotation: None, type_comment: None },
    ),
};

TypedParameter: ast::Arg = {
    <location:@L> <arg:Identifier> <a:(":" Test<"all">)?> <end_location:@R> => {
        let annotation = a.map(|x| Box::new(x.1));
        ast::Arg::new(location, end_location, ast::ArgData { arg, annotation, type_comment: None })
    },
};

StarTypedParameter: ast::Arg = {
    <location:@L> <arg:Identifier> <a:(":" TestOrStarExpr)?> <end_location:@R> => {
        let annotation = a.map(|x| Box::new(x.1));
        ast::Arg::new(location, end_location, ast::ArgData { arg, annotation, type_comment: None })
    },
};

// Use inline here to make sure the "," is not creating an ambiguity.
// TODO: figure out another grammar that makes this inline no longer required.
#[inline]
ParameterListStarArgs<ArgType, StarArgType>: (Option<Box<ast::Arg>>, Vec<ast::Arg>, Vec<ast::Expr>, Option<Box<ast::Arg>>) = {
    <location:@L> "*" <va:StarArgType?> <kw:("," ParameterDef<ArgType>)*> <kwarg:("," KwargParameter<ArgType>)?> =>? {
        // Extract keyword arguments:
        let mut kwonlyargs = Vec::new();
        let mut kw_defaults = Vec::new();
        let mut kwargs = Vec::new();
        for (name, value) in kw.into_iter().map(|x| x.1) {
            if let Some(value) = value {
                kwonlyargs.push(name);
                kw_defaults.push(value);
            } else {
                kwargs.push(name);
            }
        }
        kwargs.extend(kwonlyargs.into_iter());

        if va.is_none() && kwargs.is_empty() && kwarg.is_none() {
            Err(LexicalError {
                error: LexicalErrorType::OtherError("named arguments must follow bare *".to_string()),
                location,
            })?
        }

        let kwarg = kwarg.map(|n| n.1).flatten();
        let va = va.map(Box::new);

        Ok((va, kwargs, kw_defaults, kwarg))
    }
};

KwargParameter<ArgType>: Option<Box<ast::Arg>> = {
    "**" <kwarg:ArgType?> => {
        kwarg.map(Box::new)
    }
};

ClassDef: ast::Stmt = {
    <decorator_list:Decorator*> <location:@L> "class" <name:Identifier> <a:("(" ArgumentList ")")?> ":" <body:Suite> => {
        let (bases, keywords) = match a {
            Some((_, arg, _)) => (arg.args, arg.keywords),
            None => (vec![], vec![]),
        };
        let end_location = body.last().unwrap().end();
        ast::Stmt::new(
            location,
            end_location,
            ast::StmtKind::ClassDef {
                name,
                bases,
                keywords,
                body,
                decorator_list,
            },
        )
    },
};

// Decorators:
Decorator: ast::Expr = {
    <location:@L> "@" <p:NamedExpressionTest> "\n" => {
        p
    },
};

YieldExpr: ast::Expr = {
    <location:@L> "yield" <value:TestList?> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Yield { value: value.map(Box::new) }
    ),
    <location:@L> "yield" "from" <e:Test<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::YieldFrom { value: Box::new(e) }
    ),
};

Test<Goal>: ast::Expr = {
    <location:@L> <body:OrTest<"all">> "if" <test:OrTest<"all">> "else" <orelse:Test<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::IfExp {
            test: Box::new(test),
            body: Box::new(body),
            orelse: Box::new(orelse),
        }
    ),
    OrTest<Goal>,
    LambdaDef,
};

NamedExpressionTest: ast::Expr = {
    NamedExpression,
    Test<"all">,
}

NamedExpression: ast::Expr = {
    <location:@L> <id:Identifier> <end_location:@R> ":=" <value:Test<"all">> => {
        ast::Expr::new(
            location,
            value.end(),
            ast::ExprKind::NamedExpr {
                target: Box::new(ast::Expr::new(
                    location,
                    end_location,
                    ast::ExprKind::Name { id, ctx: ast::ExprContext::Store },
                )),
                value: Box::new(value),
            }
        )
    },
};

LambdaDef: ast::Expr = {
    <location:@L> "lambda" <p:ParameterList<UntypedParameter, UntypedParameter>?> ":" <body:Test<"all">> <end_location:@R> =>? {
        let p = validate_arguments(
            p.unwrap_or_else(|| {
                ast::Arguments {
                    posonlyargs: vec![],
                    args: vec![],
                    vararg: None,
                    kwonlyargs: vec![],
                    kw_defaults: vec![],
                    kwarg: None,
                    defaults: vec![]
                }
            }
        ))?;

        Ok(ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::Lambda {
                args: Box::new(p),
                body: Box::new(body)
            }
        ))
    }
}

OrTest<Goal>: ast::Expr = {
    <location:@L> <e1:AndTest<"all">> <e2:("or" AndTest<"all">)+> <end_location:@R> => {
        let mut values = vec![e1];
        values.extend(e2.into_iter().map(|e| e.1));
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::BoolOp { op: ast::Boolop::Or, values }
        )
    },
    AndTest<Goal>,
};

AndTest<Goal>: ast::Expr = {
    <location:@L> <e1:NotTest<"all">> <e2:("and" NotTest<"all">)+> <end_location:@R> => {
        let mut values = vec![e1];
        values.extend(e2.into_iter().map(|e| e.1));
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::BoolOp { op: ast::Boolop::And, values }
        )
    },
    NotTest<Goal>,
};

NotTest<Goal>: ast::Expr = {
    <location:@L> "not" <e:NotTest<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::UnaryOp { operand: Box::new(e), op: ast::Unaryop::Not }
    ),
    Comparison<Goal>,
};

Comparison<Goal>: ast::Expr = {
    <location:@L> <left:Expression<"all">> <comparisons:(CompOp Expression<"all">)+> <end_location:@R> => {
        let (ops, comparators) = comparisons.into_iter().unzip();
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::Compare { left: Box::new(left), ops, comparators }
        )
    },
    Expression<Goal>,
};

CompOp: ast::Cmpop = {
    "==" => ast::Cmpop::Eq,
    "!=" => ast::Cmpop::NotEq,
    "<" => ast::Cmpop::Lt,
    "<=" => ast::Cmpop::LtE,
    ">" => ast::Cmpop::Gt,
    ">=" => ast::Cmpop::GtE,
    "in" => ast::Cmpop::In,
    "not" "in" => ast::Cmpop::NotIn,
    "is" => ast::Cmpop::Is,
    "is" "not" => ast::Cmpop::IsNot,
};

Expression<Goal>: ast::Expr = {
    <location:@L> <e1:Expression<"all">> "|" <e2:XorExpression<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp { left: Box::new(e1), op: ast::Operator::BitOr, right: Box::new(e2) }
    ),
    XorExpression<Goal>,
};

XorExpression<Goal>: ast::Expr = {
    <location:@L> <e1:XorExpression<"all">> "^" <e2:AndExpression<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp { left: Box::new(e1), op: ast::Operator::BitXor, right: Box::new(e2) }
    ),
    AndExpression<Goal>,
};

AndExpression<Goal>: ast::Expr = {
    <location:@L> <e1:AndExpression<"all">> "&" <e2:ShiftExpression<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp { left: Box::new(e1), op: ast::Operator::BitAnd, right: Box::new(e2) }
    ),
    ShiftExpression<Goal>,
};

ShiftExpression<Goal>: ast::Expr = {
    <location:@L> <e1:ShiftExpression<"all">> <op:ShiftOp> <e2:ArithmeticExpression<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp { left: Box::new(e1), op, right: Box::new(e2) }
    ),
    ArithmeticExpression<Goal>,
};

ShiftOp: ast::Operator = {
    "<<" => ast::Operator::LShift,
    ">>" => ast::Operator::RShift,
};

ArithmeticExpression<Goal>: ast::Expr = {
    <location:@L> <a:ArithmeticExpression<"all">> <op:AddOp> <b:Term<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp { left: Box::new(a), op, right: Box::new(b) }
    ),
    Term<Goal>,
};

AddOp: ast::Operator = {
    "+" => ast::Operator::Add,
    "-" => ast::Operator::Sub,
};

Term<Goal>: ast::Expr = {
    <location:@L> <a:Term<"all">> <op:MulOp> <b:Factor<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp { left: Box::new(a), op, right: Box::new(b) }
    ),
    Factor<Goal>,
};

MulOp: ast::Operator = {
    "*" => ast::Operator::Mult,
    "/" => ast::Operator::Div,
    "//" => ast::Operator::FloorDiv,
    "%" => ast::Operator::Mod,
    "@" => ast::Operator::MatMult,
};

Factor<Goal>: ast::Expr = {
    <location:@L> <op:UnaryOp> <e:Factor<"all">> <end_location:@R>  => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::UnaryOp { operand: Box::new(e), op }
    ),
    Power<Goal>,
};

UnaryOp: ast::Unaryop = {
    "+" => ast::Unaryop::UAdd,
    "-" => ast::Unaryop::USub,
    "~" => ast::Unaryop::Invert,
};

Power<Goal>: ast::Expr = {
    <location:@L> <e:AtomExpr<"all">> "**" <b:Factor<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::BinOp { left: Box::new(e), op: ast::Operator::Pow, right: Box::new(b) }
    ),
    AtomExpr<Goal>,
};

AtomExpr<Goal>: ast::Expr = {
    <location:@L> "await" <atom:AtomExpr2<"all">> <end_location:@R> => {
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::Await { value: Box::new(atom) }
        )
    },
    AtomExpr2<Goal>,
}

AtomExpr2<Goal>: ast::Expr = {
    Atom<Goal>,
    <location:@L> <f:AtomExpr2<"all">> "(" <a:ArgumentList> ")" <end_location:@R> => {
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::Call { func: Box::new(f), args: a.args, keywords: a.keywords }
        )
    },
    <location:@L> <e:AtomExpr2<"all">> "[" <s:SubscriptList> "]" <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Subscript { value: Box::new(e), slice: Box::new(s), ctx: ast::ExprContext::Load }
    ),
    <location:@L> <e:AtomExpr2<"all">> "." <attr:Identifier> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Attribute { value: Box::new(e), attr, ctx: ast::ExprContext::Load }
    ),
};

SubscriptList: ast::Expr = {
    <location:@L> <s1:Subscript> <s2:("," Subscript)*> <trailing_comma:","?> <end_location:@R> => {
        if s2.is_empty() && trailing_comma.is_none() {
            s1
        } else {
            let mut dims = vec![s1];
            for x in s2 {
                dims.push(x.1)
            }

            ast::Expr::new(
                location,
                end_location,
                ast::ExprKind::Tuple { elts: dims, ctx: ast::ExprContext::Load },
            )
        }
    }
};

Subscript: ast::Expr = {
    TestOrStarNamedExpr,
    <location:@L> <e1:Test<"all">?> ":" <e2:Test<"all">?> <e3:SliceOp?> <end_location:@R> => {
        let lower = e1.map(Box::new);
        let upper = e2.map(Box::new);
        let step = e3.flatten().map(Box::new);
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::Slice { lower, upper, step }
        )
    }
};

SliceOp: Option<ast::Expr> = {
    <location:@L> ":" <e:Test<"all">?> => e,
}

Atom<Goal>: ast::Expr = {
    <location:@L> <s:(@L string @R)+> =>? Ok(parse_strings(s)?),
    <location:@L> <value:Constant> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Constant { value, kind: None }
    ),
    <location:@L> <name:Identifier> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Name { id: name, ctx: ast::ExprContext::Load }
    ),
    <location:@L> "[" <e:ListLiteralValues?> "]"<end_location:@R>  => {
        let elts = e.unwrap_or_default();
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::List { elts, ctx: ast::ExprContext::Load }
        )
    },
    <location:@L> "[" <elt:TestOrStarNamedExpr> <generators:CompFor> "]" <end_location:@R> => {
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::ListComp { elt: Box::new(elt), generators }
        )
    },
    <location:@L> "(" <elts:OneOrMore<Test<"all">>> <trailing_comma:","?> ")" <end_location:@R> if Goal != "no-withitems" => {
        if elts.len() == 1 && trailing_comma.is_none() {
            elts.into_iter().next().unwrap()
        } else {
            ast::Expr::new(
                location,
                end_location,
                ast::ExprKind::Tuple { elts, ctx: ast::ExprContext::Load },
            )
        }
    },
    <location:@L> "(" <left:(<OneOrMore<Test<"all">>> ",")?> <mid:NamedOrStarExpr> <right:("," <TestOrStarNamedExpr>)*> <trailing_comma:","?> ")" <end_location:@R> =>? {
        if left.is_none() && right.is_empty() && trailing_comma.is_none() {
            if matches!(mid.node, ast::ExprKind::Starred { .. }) {
                Err(LexicalError{
                    error: LexicalErrorType::OtherError("cannot use starred expression here".to_string()),
                    location: mid.start(),
                })?
            }
            Ok(mid)
        } else {
            let elts = left.into_iter().flatten().chain([mid]).chain(right).collect();
            Ok(ast::Expr::new(
                location,
                end_location,
                ast::ExprKind::Tuple { elts, ctx: ast::ExprContext::Load },
            ))
        }
    },
    <location:@L> "(" ")" <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Tuple { elts: Vec::new(), ctx: ast::ExprContext::Load }
    ),
    "(" <e:YieldExpr> ")" => e,
    <location:@L> "(" <elt:NamedExpressionTest> <generators:CompFor> ")" <end_location:@R> => {
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::GeneratorExp { elt: Box::new(elt), generators }
        )
    },
    "(" <location:@L> "**" <e:Expression<"all">> ")" <end_location:@R> =>? {
        Err(LexicalError{
            error : LexicalErrorType::OtherError("cannot use double starred expression here".to_string()),
            location,
        }.into())
    },
    <location:@L> "{" <e:DictLiteralValues?> "}" <end_location:@R> => {
        let (keys, values) = e
            .unwrap_or_default()
            .into_iter()
            .map(|(k, v)| (k.map(|x| *x), v))
            .unzip();
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::Dict { keys, values }
        )
    },
    <location:@L> "{" <e1:DictEntry> <generators:CompFor> "}" <end_location:@R> => {
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::DictComp {
                key: Box::new(e1.0),
                value: Box::new(e1.1),
                generators,
            }
        )
    },
    <location:@L> "{" <elts:SetLiteralValues> "}" <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Set { elts }
    ),
    <location:@L> "{" <elt:NamedExpressionTest> <generators:CompFor> "}" <end_location:@R> => {
        ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::SetComp { elt: Box::new(elt), generators }
        )
    },
    <location:@L> "True" <end_location:@R> => ast::Expr::new(location, end_location, ast::ExprKind::Constant { value: true.into(), kind: None }),
    <location:@L> "False" <end_location:@R> => ast::Expr::new(location, end_location, ast::ExprKind::Constant { value: false.into(), kind: None }),
    <location:@L> "None" <end_location:@R> => ast::Expr::new(location, end_location, ast::ExprKind::Constant { value: ast::Constant::None, kind: None }),
    <location:@L> "..." <end_location:@R> => ast::Expr::new(location, end_location, ast::ExprKind::Constant { value: ast::Constant::Ellipsis, kind: None }),
};

ListLiteralValues: Vec<ast::Expr> = {
    <e:OneOrMore<TestOrStarNamedExpr>> ","? => e,
};

DictLiteralValues: Vec<(Option<Box<ast::Expr>>, ast::Expr)> = {
    <elements:OneOrMore<DictElement>> ","? => elements,
};

DictEntry: (ast::Expr, ast::Expr) = {
    <e1: Test<"all">> ":" <e2: Test<"all">> => (e1, e2),
};

DictElement: (Option<Box<ast::Expr>>, ast::Expr) = {
    <e:DictEntry> => (Some(Box::new(e.0)), e.1),
    "**" <e:Expression<"all">> => (None, e),
};

SetLiteralValues: Vec<ast::Expr> = {
    <e1:OneOrMore<TestOrStarNamedExpr>> ","? => e1
};

ExpressionOrStarExpression = {
    Expression<"all">,
    StarExpr
};

ExpressionList: ast::Expr = {
    GenericList<ExpressionOrStarExpression>
};

ExpressionList2: Vec<ast::Expr> = {
    <elements:OneOrMore<ExpressionOrStarExpression>> ","? => elements,
};

// A test list is one of:
// - a list of expressions
// - a single expression
// - a single expression followed by a trailing comma
#[inline]
TestList: ast::Expr = {
    GenericList<TestOrStarExpr>
};

GenericList<Element>: ast::Expr = {
    <location:@L> <elts:OneOrMore<Element>> <trailing_comma:","?> <end_location:@R> => {
        if elts.len() == 1 && trailing_comma.is_none() {
            elts.into_iter().next().unwrap()
        } else {
            ast::Expr::new(
                location,
                end_location,
                ast::ExprKind::Tuple { elts, ctx: ast::ExprContext::Load }
            )
        }
    }
}

// Test
StarExpr: ast::Expr = {
    <location:@L> "*" <e:Expression<"all">> <end_location:@R> => ast::Expr::new(
        location,
        end_location,
        ast::ExprKind::Starred { value: Box::new(e), ctx: ast::ExprContext::Load },
    )
};

// Comprehensions:
CompFor: Vec<ast::Comprehension> = <c:SingleForComprehension+> => c;

SingleForComprehension: ast::Comprehension = {
    <location:@L> <is_async:"async"?> "for" <target:ExpressionList> "in" <iter:OrTest<"all">> <ifs:ComprehensionIf*> <end_location:@R> => {
        let is_async = is_async.is_some();
        ast::Comprehension {
            target: set_context(target, ast::ExprContext::Store),
            iter,
            ifs,
            is_async: if is_async { 1 } else { 0 },
        }
    }
};

ExpressionNoCond: ast::Expr = OrTest<"all">;
ComprehensionIf: ast::Expr = "if" <c:ExpressionNoCond> => c;

ArgumentList: ArgumentList = {
    <e: Comma<FunctionArgument>> =>? {
        let arg_list = parse_args(e)?;
        Ok(arg_list)
    }
};

FunctionArgument: (Option<(ast::Location, ast::Location, Option<String>)>, ast::Expr) = {
    <location:@L> <e:NamedExpressionTest> <c:CompFor?> <end_location:@R> => {
        let expr = match c {
            Some(c) => ast::Expr::new(
                location,
                end_location,
                ast::ExprKind::GeneratorExp {
                    elt: Box::new(e),
                    generators: c,
                }
            ),
            None => e,
        };
        (None, expr)
    },
    <location:@L> <i:Identifier> "=" <e:Test<"all">> <end_location:@R> => (Some((location, end_location, Some(i))), e),
    <location:@L> "*" <e:Test<"all">> <end_location:@R> => {
        let expr = ast::Expr::new(
            location,
            end_location,
            ast::ExprKind::Starred { value: Box::new(e), ctx: ast::ExprContext::Load },
        );
        (None, expr)
    },
    <location:@L> "**" <e:Test<"all">> <end_location:@R> => (Some((location, end_location, None)), e),
};

#[inline]
Comma<T>: Vec<T> = {
    <items: (<T> ",")*> <last: T?> => {
        let mut items = items;
        items.extend(last);
        items
    }
};

#[inline]
OneOrMore<T>: Vec<T> = {
    <i1: T> <i2:("," T)*> => {
        let mut items = vec![i1];
        items.extend(i2.into_iter().map(|e| e.1));
        items
    }
};

Constant: ast::Constant = {
    <value:int> => ast::Constant::Int(value),
    <value:float> => ast::Constant::Float(value),
    <s:complex> => ast::Constant::Complex { real: s.0, imag: s.1 },
};

Identifier: String = <s:name> => s;

// Hook external lexer:
extern {
    type Location = ast::Location;
    type Error = LexicalError;

    enum token::Tok {
        Indent => token::Tok::Indent,
        Dedent => token::Tok::Dedent,
        StartModule => token::Tok::StartModule,
        StartInteractive => token::Tok::StartInteractive,
        StartExpression => token::Tok::StartExpression,
        "+" => token::Tok::Plus,
        "-" => token::Tok::Minus,
        "~" => token::Tok::Tilde,
        ":" => token::Tok::Colon,
        "." => token::Tok::Dot,
        "..." => token::Tok::Ellipsis,
        "," => token::Tok::Comma,
        "*" => token::Tok::Star,
        "**" => token::Tok::DoubleStar,
        "&" => token::Tok::Amper,
        "@" => token::Tok::At,
        "%" => token::Tok::Percent,
        "//" => token::Tok::DoubleSlash,
        "^" => token::Tok::CircumFlex,
        "|" => token::Tok::Vbar,
        "<<" => token::Tok::LeftShift,
        ">>" => token::Tok::RightShift,
        "/" => token::Tok::Slash,
        "(" => token::Tok::Lpar,
        ")" => token::Tok::Rpar,
        "[" => token::Tok::Lsqb,
        "]" => token::Tok::Rsqb,
        "{" => token::Tok::Lbrace,
        "}" => token::Tok::Rbrace,
        "=" => token::Tok::Equal,
        "+=" => token::Tok::PlusEqual,
        "-=" => token::Tok::MinusEqual,
        "*=" => token::Tok::StarEqual,
        "@=" => token::Tok::AtEqual,
        "/=" => token::Tok::SlashEqual,
        "%=" => token::Tok::PercentEqual,
        "&=" => token::Tok::AmperEqual,
        "|=" => token::Tok::VbarEqual,
        "^=" => token::Tok::CircumflexEqual,
        "<<=" => token::Tok::LeftShiftEqual,
        ">>=" => token::Tok::RightShiftEqual,
        "**=" => token::Tok::DoubleStarEqual,
        "//=" => token::Tok::DoubleSlashEqual,
        ":=" => token::Tok::ColonEqual,
        "==" => token::Tok::EqEqual,
        "!=" => token::Tok::NotEqual,
        "<" => token::Tok::Less,
        "<=" => token::Tok::LessEqual,
        ">" => token::Tok::Greater,
        ">=" => token::Tok::GreaterEqual,
        "->" => token::Tok::Rarrow,
        "and" => token::Tok::And,
        "as" => token::Tok::As,
        "assert" => token::Tok::Assert,
        "async" => token::Tok::Async,
        "await" => token::Tok::Await,
        "break" => token::Tok::Break,
        "class" => token::Tok::Class,
        "continue" => token::Tok::Continue,
        "def" => token::Tok::Def,
        "del" => token::Tok::Del,
        "elif" => token::Tok::Elif,
        "else" => token::Tok::Else,
        "except" => token::Tok::Except,
        "finally" => token::Tok::Finally,
        "for" => token::Tok::For,
        "from" => token::Tok::From,
        "global" => token::Tok::Global,
        "if" => token::Tok::If,
        "import" => token::Tok::Import,
        "in" => token::Tok::In,
        "is" => token::Tok::Is,
        "lambda" => token::Tok::Lambda,
        "nonlocal" => token::Tok::Nonlocal,
        "not" => token::Tok::Not,
        "or" => token::Tok::Or,
        "pass" => token::Tok::Pass,
        "raise" => token::Tok::Raise,
        "return" => token::Tok::Return,
        "try" => token::Tok::Try,
        "while" => token::Tok::While,
        "match" => token::Tok::Match,
        "case" => token::Tok::Case,
        "with" => token::Tok::With,
        "yield" => token::Tok::Yield,
        "True" => token::Tok::True,
        "False" => token::Tok::False,
        "None" => token::Tok::None,
        int => token::Tok::Int { value: <BigInt> },
        float => token::Tok::Float { value: <f64> },
        complex => token::Tok::Complex { real: <f64>, imag: <f64> },
        string => token::Tok::String {
            value: <String>,
            kind: <StringKind>,
            triple_quoted: <bool>
        },
        name => token::Tok::Name { name: <String> },
        "\n" => token::Tok::Newline,
        ";" => token::Tok::Semi,
        "#" => token::Tok::Comment(_),
    }
}