Grcov report - comprehension.rs

1

use crate::RecoverableParseError;

2

use crate::errors::FatalParseError;

3

use crate::expression::parse_expression;

4

use crate::field;

5

use crate::parser::ParseContext;

6

use crate::parser::domain::parse_domain;

7

use crate::util::{TypecheckingContext, named_children};

8

use conjure_cp_core::ast::ac_operators::ACOperatorKind;

9

use conjure_cp_core::ast::comprehension::ComprehensionBuilder;

10

use conjure_cp_core::ast::serde::HasId as _;

11

use conjure_cp_core::ast::{

12

    Atom, DeclarationPtr, Expression, Literal, Metadata, Moo, Name, Reference,

13

};

14

use tree_sitter::Node;

15

use uniplate::Uniplate as _;

16

17

/// Collects the names a generator pattern binds, each with its index path into the generator

18

/// variable: `(i, _)` binds `i` at path `[1]`, and `(i, (j, k))` binds `k` at path `[2, 2]`.

19

///

20

/// Patterns are written as tuples, so an element arrives wrapped in whatever expression nodes a

21

/// tuple element goes through; only the identifiers at the bottom matter here.

22

///

23

/// Returns false for a pattern this cannot read, having recorded an error.

24

fn collect_pattern_bindings(

25

    ctx: &mut ParseContext,

26

    node: &Node,

27

    path: &mut Vec<i32>,

28

    bindings: &mut Vec<(Name, Vec<i32>)>,

29

) -> bool {

30

    match node.kind() {

31

        "identifier" => {

32

            let name = &ctx.source_code[node.start_byte()..node.end_byte()];

33

            // `_` names a component the pattern discards.

34

            if name != "_" {

35

                bindings.push((Name::user(name), path.clone()));

36

37

            true

38

39

        "tuple" => {

40

            for (position, element) in named_children(node).enumerate() {

41

                path.push(position as i32 + 1);

42

                let bound = collect_pattern_bindings(ctx, &element, path, bindings);

43

                path.pop();

44

                if !bound {

45

                    return false;

46

47

48

            true

49

50

        // A tuple element is wrapped in expression nodes; unwrap to what it actually holds.

51

        _ => match named_children(node).collect::<Vec<_>>().as_slice() {

52

            [inner] => collect_pattern_bindings(ctx, inner, path, bindings),

53

            _ => {

54

                ctx.record_error(RecoverableParseError::new(

55

                    format!(

56

                        "Expected a name or `_` in a generator pattern, got '{}'",

57

                        &ctx.source_code[node.start_byte()..node.end_byte()]

58

),

59

                    Some(node.range()),

60

));

61

                false

62

63

},

64

65

66

67

/// Binds the names a generator pattern introduces, for a generator already added to `builder`

68

/// under `var_name`.

69

///

70

/// Patterns are handled exactly like comprehension lettings: each name is bound to a projection

71

/// out of the generator variable and substituted into everything that follows, so nothing

72

/// downstream has to know that a pattern was written.

73

268809

fn bind_generator_pattern(

74

268809

    ctx: &mut ParseContext,

75

268809

    builder: &mut ComprehensionBuilder,

76

268809

    var_node: &Node,

77

268809

    var_name: &Name,

78

268809

    lettings: &mut Vec<(DeclarationPtr, Expression)>,

79

268809

) -> bool {

80

268809

    if var_node.kind() != "tuple" {

81

268809

        return true;

82

83

84

    let mut bindings = Vec::new();

85

    if !collect_pattern_bindings(ctx, var_node, &mut Vec::new(), &mut bindings) {

86

        return false;

87

88

89

    let symbols = builder.generator_symboltable();

90

    let Some(generator_decl) = symbols.read().lookup_local(var_name) else {

91

        ctx.record_error(RecoverableParseError::new(

92

            format!("Generator pattern variable {var_name} is not in scope"),

93

            Some(var_node.range()),

94

));

95

        return false;

96

};

97

98

    for (name, path) in bindings {

99

        let mut value = Expression::from(Reference::new(generator_decl.clone()));

100

        for index in path {

101

            value = Expression::SafeIndex(

102

                Metadata::new(),

103

                Moo::new(value),

104

                vec![Expression::from(Literal::Int(index))],

105

);

106

107

        let decl = DeclarationPtr::new_value_letting(name, value.clone());

108

        symbols.write().insert(decl.clone());

109

        lettings.push((decl, value));

110

111

112

    true

113

268809

114

115

63808

fn parse_generator(

116

63808

    ctx: &mut ParseContext,

117

63808

    mut builder: ComprehensionBuilder,

118

63808

    generator_node: &Node,

119

63808

    lettings: &mut Vec<(DeclarationPtr, Expression)>,

120

63808

) -> Result<Option<ComprehensionBuilder>, FatalParseError> {

121

63808

    let Some(var_node) = field!(recover, ctx, generator_node, "variable") else {

122

        return Ok(None);

123

};

124

63808

    let var_name_str = &ctx.source_code[var_node.start_byte()..var_node.end_byte()];

125

63808

    let var_name = Name::user(var_name_str);

126

127

63808

    if let Some(domain_node) = generator_node.child_by_field_name("domain") {

128

43492

        let mut domain_ctx = ctx.with_new_symbols(Some(builder.generator_symboltable()));

129

43492

        let Some(var_domain) = parse_domain(&mut domain_ctx, domain_node)? else {

130

24

            return Ok(None);

131

};

132

43468

        let decl = DeclarationPtr::new_find(var_name.clone(), var_domain);

133

43468

        let mut builder = builder.generator(decl);

134

43468

        if !bind_generator_pattern(ctx, &mut builder, &var_node, &var_name, lettings) {

135

            return Ok(None);

136

43468

137

43468

        return Ok(Some(builder));

138

20316

139

140

20316

    if let Some(collection_node) = generator_node.child_by_field_name("collection") {

141

20316

        let mut collection_ctx = ctx.with_new_symbols(Some(builder.generator_symboltable()));

142

20316

        collection_ctx.typechecking_context = TypecheckingContext::Unknown;

143

20316

        collection_ctx.inner_typechecking_context = TypecheckingContext::Unknown;

144

20316

        let Some(collection_expr) = parse_expression(&mut collection_ctx, collection_node)? else {

145

            return Ok(None);

146

};

147

20316

        let collection_expr = substitute_lettings(collection_expr, lettings);

148

20316

        let mut builder = builder.expression_generator(var_name.clone(), collection_expr);

149

20316

        if !bind_generator_pattern(ctx, &mut builder, &var_node, &var_name, lettings) {

150

            return Ok(None);

151

20316

152

20316

        return Ok(Some(builder));

153

154

155

    ctx.record_error(RecoverableParseError::new(

156

        format!(

157

            "Generator requires either a domain or collection (e.g. `{var_name} : domain` or `{var_name} <- expr`)"

158

),

159

        Some(generator_node.range()),

160

));

161

    Ok(None)

162

63808

163

164

179729

fn parse_quantifier_variables<'a>(

165

179729

    ctx: &mut ParseContext,

166

179729

    node: &Node<'a>,

167

179729

) -> Option<Vec<(Name, Node<'a>)>> {

168

179729

    let mut variables = Vec::new();

169

179729

    let mut cursor = node.walk();

170

179729

    for child in node.children_by_field_name("variables", &mut cursor) {

171

230321

        if matches!(child.kind(), "identifier" | "tuple") {

172

205025

            let var_name_str = &ctx.source_code[child.start_byte()..child.end_byte()];

173

205025

            variables.push((Name::user(var_name_str), child));

174

25296

175

176

177

179729

    if variables.is_empty() {

178

        ctx.record_error(RecoverableParseError::new(

179

            "Quantifier and aggregate expressions require variables".to_string(),

180

            Some(node.range()),

181

));

182

        return None;

183

179729

184

185

179729

    Some(variables)

186

179729

187

188

27552

fn add_quantifier_generators_from_collection(

189

27552

    ctx: &mut ParseContext,

190

27552

    mut builder: ComprehensionBuilder,

191

27552

    variables: &[(Name, Node)],

192

27552

    collection_node: Node,

193

27552

    lettings: &mut Vec<(DeclarationPtr, Expression)>,

194

27552

) -> Result<Option<ComprehensionBuilder>, FatalParseError> {

195

27552

    let mut collection_ctx = ctx.with_new_symbols(Some(builder.generator_symboltable()));

196

27552

    collection_ctx.typechecking_context = TypecheckingContext::Unknown;

197

27552

    collection_ctx.inner_typechecking_context = TypecheckingContext::Unknown;

198

27552

    let Some(collection_expr) = parse_expression(&mut collection_ctx, collection_node)? else {

199

        return Ok(None);

200

};

201

202

27552

    for (var_name, var_node) in variables {

203

27552

        builder = builder.expression_generator(var_name.clone(), collection_expr.clone());

204

27552

        if !bind_generator_pattern(ctx, &mut builder, var_node, var_name, lettings) {

205

            return Ok(None);

206

27552

207

208

209

27552

    Ok(Some(builder))

210

27552

211

212

/// Replaces references to comprehension lettings with the expressions they are bound to.

213

///

214

/// Comprehension lettings are pure sugar. They cannot be left as declarations for the rewriter:

215

/// their value usually mentions the generator variables, which comprehension expansion substitutes

216

/// in the comprehension body but not in declarations stored in the symbol table.

217

296129

fn substitute_lettings(expr: Expression, lettings: &[(DeclarationPtr, Expression)]) -> Expression {

218

296129

    if lettings.is_empty() {

219

296125

        return expr;

220

4

221

222

12

    expr.transform(&|expr| match &expr {

223

8

        Expression::Atomic(_, Atom::Reference(reference)) => lettings

224

8

            .iter()

225

8

            .find(|(decl, _)| decl.id() == reference.id())

226

4

            .map(|(_, value)| value.clone())

227

8

            .unwrap_or(expr),

228

4

        _ => expr,

229

12

})

230

296129

231

232

61592

pub fn parse_comprehension(

233

61592

    ctx: &mut ParseContext,

234

61592

    node: &Node,

235

61592

) -> Result<Option<Expression>, FatalParseError> {

236

    // If we're in a set context, add error and return early since comprehensions don't produce sets

237

61592

    if ctx.typechecking_context == crate::util::TypecheckingContext::Set {

238

        ctx.record_error(crate::errors::RecoverableParseError::new(

239

            format!(

240

                "Type error: {}\n\tExpected: set\n\tGot: comprehension",

241

                ctx.source_code[node.start_byte()..node.end_byte()].trim()

242

),

243

            Some(node.range()),

244

));

245

61592

246

247

    // Comprehensions require a symbol table passed in

248

61592

    let symbols_ptr = match ctx.symbols.clone() {

249

61592

        Some(s) => s,

250

        None => {

251

            ctx.record_error(RecoverableParseError::new(

252

                "Comprehensions require a symbol table".to_string(),

253

                Some(node.range()),

254

));

255

            return Ok(None);

256

257

};

258

259

61592

    let mut builder = ComprehensionBuilder::new(symbols_ptr);

260

261

    // We need to track the return expression node separately since it appears first in syntax

262

    // but we need to parse generators first (to get variables in scope)

263

61592

    let mut return_expr_node: Option<Node> = None;

264

265

    // Names bound by `letting` qualifiers, substituted back into everything that follows.

266

61592

    let mut lettings: Vec<(DeclarationPtr, Expression)> = Vec::new();

267

268

    // set return expression node and parse generators/conditions

269

61592

    for child in named_children(node) {

270

160012

        match child.kind() {

271

            "arithmetic_expr" | "bool_expr" | "comparison_expr" | "atom" => {

272

                // Store the return expression node to parse later

273

61592

                return_expr_node = Some(child);

274

275

            "generator" => {

276

63808

                let Some(updated_builder) = parse_generator(ctx, builder, &child, &mut lettings)?

277

                else {

278

24

                    return Ok(None);

279

};

280

63784

                builder = updated_builder;

281

282

            "condition" => {

283

                // Parse the condition expression

284

34512

                let Some(expr_node) = field!(recover, ctx, child, "expression") else {

285

                    return Ok(None);

286

};

287

34512

                let generator_symboltable = builder.generator_symboltable();

288

289

                // Parse with a new context using the generator symbol table

290

34512

                let mut guard_ctx = ctx.with_new_symbols(Some(generator_symboltable));

291

34512

                let Some(guard_expr) = parse_expression(&mut guard_ctx, expr_node)? else {

292

                    return Ok(None);

293

};

294

295

                // Add the condition as a guard

296

34512

                builder = builder.guard(substitute_lettings(guard_expr, &lettings));

297

298

            "comprehension_letting" => {

299

4

                let Some(name_node) = field!(recover, ctx, child, "name") else {

300

                    return Ok(None);

301

};

302

4

                let Some(value_node) = field!(recover, ctx, child, "value") else {

303

                    return Ok(None);

304

};

305

306

4

                let generator_symboltable = builder.generator_symboltable();

307

4

                let mut letting_ctx = ctx.with_new_symbols(Some(generator_symboltable.clone()));

308

4

                letting_ctx.typechecking_context = TypecheckingContext::Unknown;

309

4

                letting_ctx.inner_typechecking_context = TypecheckingContext::Unknown;

310

4

                let Some(value) = parse_expression(&mut letting_ctx, value_node)? else {

311

                    return Ok(None);

312

};

313

4

                let value = substitute_lettings(value, &lettings);

314

315

                let name =

316

4

                    Name::user(&ctx.source_code[name_node.start_byte()..name_node.end_byte()]);

317

                // The declaration only exists so that the name resolves while parsing what

318

                // follows; every reference to it is substituted away before the comprehension is

319

                // built.

320

4

                let decl = DeclarationPtr::new_value_letting(name, value.clone());

321

4

                generator_symboltable.write().insert(decl.clone());

322

4

                lettings.push((decl, value));

323

324

96

            _ => {

325

                // Skip other nodes (like punctuation)

326

96

327

328

329

330

    // parse the return expression

331

61568

    let return_expr_node = match return_expr_node {

332

61568

        Some(node) => node,

333

        None => {

334

            ctx.record_error(RecoverableParseError::new(

335

                "Comprehension missing return expression".to_string(),

336

                Some(node.range()),

337

));

338

            return Ok(None);

339

340

};

341

342

    // Use the return expression symbol table which already has quantified variables (as Given) and parent as parent

343

    // Parse using the inner typechecking context

344

61568

    let saved_inner_ctx = ctx.inner_typechecking_context;

345

61568

    let mut return_ctx = ctx.with_new_symbols(Some(builder.return_expr_symboltable()));

346

61568

    return_ctx.typechecking_context = saved_inner_ctx;

347

61568

    return_ctx.inner_typechecking_context = TypecheckingContext::Unknown;

348

61568

    let Some(return_expr) = parse_expression(&mut return_ctx, return_expr_node)? else {

349

        return Ok(None);

350

};

351

352

    // Build the comprehension with the return expression

353

61568

    let comprehension = builder.with_return_value(substitute_lettings(return_expr, &lettings));

354

355

61568

    Ok(Some(Expression::Comprehension(

356

61568

        Metadata::new(),

357

61568

        Moo::new(comprehension),

358

61568

)))

359

61592

360

361

/// Parse comprehension-style expressions

362

/// - `forAll vars : domain . expr` → `And(Comprehension(...))`

363

/// - `sum vars : domain . expr` → `Sum(Comprehension(...))`

364

179729

pub fn parse_quantifier_or_aggregate_expr(

365

179729

    ctx: &mut ParseContext,

366

179729

    node: &Node,

367

179729

) -> Result<Option<Expression>, FatalParseError> {

368

    // Quantifier and aggregate expressions require a symbol table

369

179729

    let symbols_ptr = match ctx.symbols.clone() {

370

179729

        Some(s) => s,

371

        None => {

372

            ctx.record_error(RecoverableParseError::new(

373

                "Quantifier and aggregate expressions require a symbol table".to_string(),

374

                Some(node.range()),

375

));

376

            return Ok(None);

377

378

};

379

380

    // Create the comprehension builder

381

179729

    let mut builder = ComprehensionBuilder::new(symbols_ptr);

382

383

    // Names bound by generator patterns, substituted back into everything that follows.

384

179729

    let mut lettings: Vec<(DeclarationPtr, Expression)> = Vec::new();

385

386

179729

    let Some(variables) = parse_quantifier_variables(ctx, node) else {

387

        return Ok(None);

388

};

389

390

179729

    let domain_node = node.child_by_field_name("domain");

391

179729

    let collection_node = node.child_by_field_name("collection");

392

393

179729

    if domain_node.is_none() && collection_node.is_none() {

394

        ctx.record_error(RecoverableParseError::new(

395

            "Quantifier and aggregate expressions require a domain or collection".to_string(),

396

            Some(node.range()),

397

));

398

        return Ok(None);

399

179729

400

401

179729

    if domain_node.is_some() && collection_node.is_some() {

402

        ctx.record_error(RecoverableParseError::new(

403

            "Quantifier and aggregate expressions cannot have both a domain and a collection"

404

                .to_string(),

405

            Some(node.range()),

406

));

407

        return Ok(None);

408

179729

409

410

    // Get the operator type

411

179729

    let Some(operator_node) = field!(recover, ctx, node, "operator") else {

412

        return Ok(None);

413

};

414

179729

    let operator_str = &ctx.source_code[operator_node.start_byte()..operator_node.end_byte()];

415

416

179729

    let (ac_operator_kind, wrapper) = match operator_str {

417

101369

        "forAll" => (ACOperatorKind::And, "And"),

418

50304

        "exists" => (ACOperatorKind::Or, "Or"),

419

28056

        "sum" => (ACOperatorKind::Sum, "Sum"),

420

        // min/max are not true AC operators, but still need their own skip-operator tag so a

421

        // symbolic guard lowers correctly instead of silently substituting Sum's identity (0).

422

        "min" => (ACOperatorKind::Min, "Min"),

423

        "max" => (ACOperatorKind::Max, "Max"),

424

        _ => {

425

            ctx.record_error(RecoverableParseError::new(

426

                format!("Unknown operator: {}", operator_str),

427

                Some(operator_node.range()),

428

));

429

            return Ok(None);

430

431

};

432

433

    // Add variables as generators

434

179729

    if let Some(domain_node) = domain_node {

435

152177

        let saved_ctx = ctx.typechecking_context;

436

152177

        let saved_inner_ctx = ctx.inner_typechecking_context;

437

152177

        ctx.typechecking_context = TypecheckingContext::Unknown;

438

152177

        ctx.inner_typechecking_context = TypecheckingContext::Unknown;

439

440

152177

        let Some(domain) = parse_domain(ctx, domain_node)? else {

441

            ctx.typechecking_context = saved_ctx;

442

            ctx.inner_typechecking_context = saved_inner_ctx;

443

            return Ok(None);

444

};

445

446

152177

        ctx.typechecking_context = saved_ctx;

447

152177

        ctx.inner_typechecking_context = saved_inner_ctx;

448

449

152177

        for (var_name, var_node) in &variables {

450

177473

            let decl = DeclarationPtr::new_find(var_name.clone(), domain.clone());

451

177473

            builder = builder.generator(decl);

452

177473

            if !bind_generator_pattern(ctx, &mut builder, var_node, var_name, &mut lettings) {

453

                return Ok(None);

454

177473

455

456

27552

    } else if let Some(collection_node) = collection_node {

457

27552

        let Some(updated_builder) = add_quantifier_generators_from_collection(

458

27552

            ctx,

459

27552

            builder,

460

27552

            &variables,

461

27552

            collection_node,

462

27552

            &mut lettings,

463

27552

)?

464

        else {

465

            return Ok(None);

466

};

467

27552

        builder = updated_builder;

468

469

470

    // A guard written between the generator and the `.`, e.g. `forAll (i, _) in s, i > 1 . ...`.

471

179729

    if let Some(guard_node) = node.child_by_field_name("guard") {

472

        let mut guard_ctx = ctx.with_new_symbols(Some(builder.generator_symboltable()));

473

        guard_ctx.typechecking_context = TypecheckingContext::Unknown;

474

        guard_ctx.inner_typechecking_context = TypecheckingContext::Unknown;

475

        let Some(guard) = parse_expression(&mut guard_ctx, guard_node)? else {

476

            return Ok(None);

477

};

478

        builder = builder.guard(substitute_lettings(guard, &lettings));

479

179729

480

481

    // Parse the expression (after variables are in the symbol table)

482

179729

    let Some(expression_node) = field!(recover, ctx, node, "expression") else {

483

        return Ok(None);

484

};

485

486

    // Parse with a new context using the return expression symbol table

487

    // Prase using the inner typechecking context

488

179729

    let saved_inner_ctx = ctx.inner_typechecking_context;

489

179729

    let mut expr_ctx = ctx.with_new_symbols(Some(builder.return_expr_symboltable()));

490

179729

    expr_ctx.typechecking_context = saved_inner_ctx;

491

179729

    expr_ctx.inner_typechecking_context = TypecheckingContext::Unknown;

492

179729

    let Some(expression) = parse_expression(&mut expr_ctx, expression_node)? else {

493

        return Ok(None);

494

};

495

496

    // Build the comprehension

497

179729

    let mut comprehension = builder.with_return_value(substitute_lettings(expression, &lettings));

498

179729

    comprehension.skip_operator = Some(ac_operator_kind);

499

179729

    let wrapped_comprehension = Expression::Comprehension(Metadata::new(), Moo::new(comprehension));

500

501

    // Wrap in the appropriate expression type

502

179729

    match wrapper {

503

101369

        "And" => Ok(Some(Expression::And(

504

101369

            Metadata::new(),

505

101369

            Moo::new(wrapped_comprehension),

506

101369

        ))),

507

50304

        "Or" => Ok(Some(Expression::Or(

508

50304

            Metadata::new(),

509

50304

            Moo::new(wrapped_comprehension),

510

50304

        ))),

511

28056

        "Sum" => Ok(Some(Expression::Sum(

512

28056

            Metadata::new(),

513

28056

            Moo::new(wrapped_comprehension),

514

28056

        ))),

515

        "Min" => Ok(Some(Expression::Min(

516

            Metadata::new(),

517

            Moo::new(wrapped_comprehension),

518

        ))),

519

        "Max" => Ok(Some(Expression::Max(

520

            Metadata::new(),

521

            Moo::new(wrapped_comprehension),

522

        ))),

523

        _ => unreachable!(),

524

525

179729