1use std::path::Path;
6
7use ahash::AHashMap;
8use logos::{Lexer, Logos};
9use malachite_bigint::BigInt;
10use num_traits::{ToPrimitive, Zero};
11use veripb_formula::prelude::*;
12
13use crate::{error::ParserError, opb_token::OPBToken, parser::get_lines};
14
15type OptionalConstraint = Option<PBConstraintEnum>;
16type Constraint = PBConstraintEnum;
17
18#[derive(Debug, PartialEq, Clone, Copy)]
20enum Comparator {
21 GreaterEqual,
22 LessEqual,
23 Equal,
24}
25
26#[inline]
30pub fn parse_opb_from_file<P>(
31 filename: P,
32) -> Result<(Formula, VarNameManager, AHashMap<String, isize>), ParserError>
33where
34 P: AsRef<Path>,
35{
36 let mut var_name_manager = VarNameManager::default();
37 let (formula, labels) = parse_opb_from_file_given_var_manager(filename, &mut var_name_manager)?;
38 Ok((formula, var_name_manager, labels))
39}
40
41pub fn parse_opb_from_file_given_var_manager<P>(
45 filename: P,
46 var_name_manager: &mut VarNameManager,
47) -> Result<(Formula, AHashMap<String, isize>), ParserError>
48where
49 P: AsRef<Path>,
50{
51 let mut formula = Formula::default();
52 let mut labels = AHashMap::new();
53 let lines = get_lines(&filename)?;
54
55 for (line_number, line) in lines.map_while(Result::ok).enumerate() {
56 let mut lex = OPBToken::lexer(&line);
57 let result = match lex.next() {
58 None | Some(Ok(OPBToken::Comment)) => continue,
59 Some(Ok(OPBToken::Label)) => {
60 let label = lex.slice();
61 let result = match lex.next() {
62 Some(Ok(OPBToken::Integer)) => Some(
63 parse_opb_constraint_dyn(&mut lex, var_name_manager).map_err(|e| {
64 e.add_file_and_line(
65 filename.as_ref().to_string_lossy().to_string(),
66 line_number,
67 )
68 .unwrap()
69 })?,
70 ),
71 Some(Ok(OPBToken::GreaterEqual)) => Some(
72 parse_constraint_empty_terms(&mut lex, Comparator::GreaterEqual).map_err(
73 |e| {
74 e.add_file_and_line(
75 filename.as_ref().to_string_lossy().to_string(),
76 line_number,
77 )
78 .unwrap()
79 },
80 )?,
81 ),
82 Some(Err(_)) | Some(Ok(_)) | None => {
83 return Err(ParserError::token_error_with_file(
84 lex.span(),
85 "'>=' or integer",
86 filename.as_ref().to_string_lossy().to_string(),
87 line_number,
88 ));
89 }
90 };
91 if result.as_ref().unwrap().1.is_some() {
92 return Err(ParserError::token_error_with_file(
93 lex.span(),
94 "inequality constraint",
95 filename.as_ref().to_string_lossy().to_string(),
96 line_number,
97 ));
98 }
99 labels.insert(label.to_string(), (formula.len() + 1) as isize);
100 result
101 }
102 Some(Ok(OPBToken::Minimize)) => {
103 formula.objective = Some(parse_opb_objective(&mut lex, var_name_manager, false)?);
104 None
105 }
106 Some(Ok(OPBToken::Maximize)) => {
107 formula.objective = Some(parse_opb_objective(&mut lex, var_name_manager, true)?);
108 None
109 }
110 Some(Ok(OPBToken::Integer)) => Some(
111 parse_opb_constraint_dyn(&mut lex, var_name_manager).map_err(|e| {
112 e.add_file_and_line(
113 filename.as_ref().to_string_lossy().to_string(),
114 line_number,
115 )
116 .unwrap()
117 })?,
118 ),
119 Some(Ok(OPBToken::GreaterEqual)) => Some(
120 parse_constraint_empty_terms(&mut lex, Comparator::GreaterEqual).map_err(|e| {
121 e.add_file_and_line(
122 filename.as_ref().to_string_lossy().to_string(),
123 line_number,
124 )
125 .unwrap()
126 })?,
127 ),
128 Some(Ok(OPBToken::LessEqual)) => Some(
129 parse_constraint_empty_terms(&mut lex, Comparator::LessEqual).map_err(|e| {
130 e.add_file_and_line(
131 filename.as_ref().to_string_lossy().to_string(),
132 line_number,
133 )
134 .unwrap()
135 })?,
136 ),
137 Some(Ok(OPBToken::Equal)) => Some(
138 parse_constraint_empty_terms(&mut lex, Comparator::Equal).map_err(|e| {
139 e.add_file_and_line(
140 filename.as_ref().to_string_lossy().to_string(),
141 line_number,
142 )
143 .unwrap()
144 })?,
145 ),
146 Some(Err(_)) | Some(Ok(_)) => {
147 return Err(ParserError::token_error_with_file(
148 lex.span(),
149 "'*', 'min:', '>=', '<=', '=', or integer",
150 filename.as_ref().to_string_lossy().to_string(),
151 line_number,
152 ));
153 }
154 };
155 if let Some((geq_constraint, leq_constraint)) = result {
156 formula.constraints.push(geq_constraint);
157 if let Some(constraint) = leq_constraint {
158 formula.constraints.push(constraint);
159 }
160 }
161 }
162
163 Ok((formula, labels))
164}
165
166pub fn parse_opb_objective(
168 lex: &mut Lexer<OPBToken>,
169 var_name_manager: &mut VarNameManager,
170 is_maximization: bool,
171) -> Result<PBObjective, ParserError> {
172 let mut terms = Vec::new();
173 let mut integer = BigInt::zero();
174 match lex.next() {
175 Some(Ok(OPBToken::Integer)) => integer = lex.slice().parse().unwrap(),
176 Some(Ok(OPBToken::Semicolon)) | None => {
177 return Ok(PBObjective::from_terms(terms, integer, is_maximization))
178 }
179 _ => return Err(ParserError::token_error(lex.span(), "integer or ';'")),
180 }
181
182 while let Some(token) = lex.next() {
183 let lit = match token {
185 Ok(OPBToken::Var) => Lit::from_var(var_name_manager.add_by_name(lex.slice()), false),
186 Ok(OPBToken::Negation) => {
187 if lex.next() == Some(Ok(OPBToken::Var)) {
188 Lit::from_var(var_name_manager.add_by_name(lex.slice()), true)
189 } else {
190 return Err(ParserError::token_error(lex.span(), "variable name"));
191 }
192 }
193 Ok(OPBToken::Semicolon) => {
194 return Ok(PBObjective::from_terms(terms, integer, is_maximization));
195 }
196 _ => return Err(ParserError::token_error(lex.span(), "literal or ';'")),
197 };
198
199 terms.push(GeneralPBTerm::new(integer, lit));
200
201 match lex.next() {
203 Some(Ok(OPBToken::Integer)) => integer = lex.slice().parse().unwrap(),
204 Some(Ok(OPBToken::Semicolon)) | None => {
205 return Ok(PBObjective::from_terms(
206 terms,
207 BigInt::zero(),
208 is_maximization,
209 ))
210 }
211 _ => return Err(ParserError::token_error(lex.span(), "integer or ';'")),
212 }
213 }
214
215 Ok(PBObjective::from_terms(terms, integer, is_maximization))
216}
217
218pub fn parse_single_constraint(
220 lex: &mut Lexer<OPBToken>,
221 var_names: &mut VarNameManager,
222) -> Result<(Constraint, OptionalConstraint), ParserError> {
223 match lex.next() {
224 Some(Ok(OPBToken::Integer)) => parse_opb_constraint_dyn(lex, var_names),
225 Some(Ok(OPBToken::GreaterEqual)) => {
226 parse_constraint_empty_terms(lex, Comparator::GreaterEqual)
227 }
228 Some(Ok(OPBToken::Equal)) => parse_constraint_empty_terms(lex, Comparator::Equal),
229 _ => Err(ParserError::token_error(
230 lex.span(),
231 "'*', 'min:', '>=', '<=', or integer",
232 )),
233 }
234}
235
236fn parse_constraint_empty_terms(
238 lex: &mut Lexer<OPBToken>,
239 comparator: Comparator,
240) -> Result<(Constraint, OptionalConstraint), ParserError> {
241 is_integer(lex)?;
242 let degree: BigInt = lex.slice().parse().unwrap();
243 check_constraint_end(lex)?;
244
245 match degree.to_i64() {
246 Some(degree) => Ok(get_constraints_from_terms(comparator, vec![], 0, degree)),
247 None => match degree.to_i128() {
248 Some(degree) => Ok(get_constraints_from_terms(comparator, vec![], 0, degree)),
249 None => Ok(get_constraints_from_terms(
250 comparator,
251 vec![],
252 BigInt::zero(),
253 degree,
254 )),
255 },
256 }
257}
258
259fn parse_opb_constraint_dyn(
265 lex: &mut Lexer<OPBToken>,
266 var_name_manager: &mut VarNameManager,
267) -> Result<(Constraint, OptionalConstraint), ParserError> {
268 let mut terms_i64 = Vec::new();
270 let mut coeff_sum_i64: i64 = 0;
271 let mut comparator =
272 parse_opb_terms(lex, var_name_manager, &mut coeff_sum_i64, &mut terms_i64)?;
273
274 if let Some(comparator) = comparator {
275 if let Ok(degree) = lex.slice().parse() {
276 check_constraint_end(lex)?;
277 return Ok(get_constraints_from_terms(
278 comparator,
279 terms_i64,
280 coeff_sum_i64,
281 degree,
282 ));
283 }
284 }
285
286 let mut terms_i128 = terms_i64.into_iter().map(|t| t.into()).collect();
288 let mut coeff_sum_i128: i128 = coeff_sum_i64.into();
289 if comparator.is_none() {
290 comparator = parse_opb_terms(lex, var_name_manager, &mut coeff_sum_i128, &mut terms_i128)?;
291 }
292 if let Some(comparator) = comparator {
293 if let Ok(degree) = lex.slice().parse() {
294 check_constraint_end(lex)?;
295 return Ok(get_constraints_from_terms(
296 comparator,
297 terms_i128,
298 coeff_sum_i128,
299 degree,
300 ));
301 }
302 }
303
304 let mut terms_big: Vec<GeneralPBTerm<BigInt>> =
306 terms_i128.into_iter().map(|t| t.into()).collect();
307 let mut coeff_sum_big: BigInt = coeff_sum_i128.into();
308 let comparator = if let Some(comparator) = comparator {
309 comparator
310 } else {
311 parse_opb_terms(lex, var_name_manager, &mut coeff_sum_big, &mut terms_big)?.unwrap()
312 };
313 let degree = lex.slice().parse().unwrap();
314 check_constraint_end(lex)?;
315 Ok(get_constraints_from_terms(
316 comparator,
317 terms_big,
318 coeff_sum_big,
319 degree,
320 ))
321}
322
323#[inline]
329fn parse_opb_terms<N: Int>(
330 lex: &mut Lexer<OPBToken>,
331 var_name_manager: &mut VarNameManager,
332 coeff_sum: &mut N,
333 terms: &mut Vec<GeneralPBTerm<N>>,
334) -> Result<Option<Comparator>, ParserError> {
335 let mut coeff: N = match lex.slice().parse() {
336 Ok(integer) => integer,
337 Err(_) => return Ok(None),
338 };
339
340 loop {
341 if coeff.is_negative() {
342 let abs = match N::zero().checked_sub(&coeff) {
343 Some(abs) => abs,
344 None => return Ok(None),
345 };
346 *coeff_sum = match coeff_sum.checked_add(&abs) {
347 Some(sum) => sum,
348 None => return Ok(None),
349 };
350 } else {
351 *coeff_sum = match coeff_sum.checked_add(&coeff) {
352 Some(sum) => sum,
353 None => return Ok(None),
354 };
355 }
356 let lit = match lex.next() {
357 Some(Ok(OPBToken::Var)) => {
358 Lit::from_var(var_name_manager.add_by_name(lex.slice()), false)
359 }
360 Some(Ok(OPBToken::Negation)) => {
361 if lex.next() == Some(Ok(OPBToken::Var)) {
362 Lit::from_var(var_name_manager.add_by_name(lex.slice()), true)
363 } else {
364 return Err(ParserError::token_error(lex.span(), "variable name"));
365 }
366 }
367 _ => {
368 return Err(ParserError::token_error(lex.span(), "literal"));
369 }
370 };
371 terms.push(GeneralPBTerm::new(coeff, lit));
372
373 coeff = match lex.next() {
374 Some(Ok(OPBToken::GreaterEqual)) => {
375 is_integer(lex)?;
376 return Ok(Some(Comparator::GreaterEqual));
377 }
378 Some(Ok(OPBToken::LessEqual)) => {
379 is_integer(lex)?;
380 return Ok(Some(Comparator::LessEqual));
381 }
382 Some(Ok(OPBToken::Equal)) => {
383 is_integer(lex)?;
384 return Ok(Some(Comparator::Equal));
385 }
386 Some(Ok(OPBToken::Integer)) => match lex.slice().parse() {
387 Ok(integer) => integer,
388 Err(_) => return Ok(None),
389 },
390 _ => {
391 return Err(ParserError::token_error(
392 lex.span(),
393 "'>=', '<=', or integer",
394 ));
395 }
396 };
397 }
398}
399
400#[inline]
404fn is_integer(lex: &mut Lexer<OPBToken>) -> Result<(), ParserError> {
405 if lex.next() != Some(Ok(OPBToken::Integer)) {
406 return Err(ParserError::token_error(lex.span(), "integer"));
407 }
408 Ok(())
409}
410
411#[inline]
413fn check_constraint_end(lex: &mut Lexer<OPBToken>) -> Result<(), ParserError> {
414 if lex.next() != Some(Ok(OPBToken::Semicolon)) {
415 Err(ParserError::token_error(lex.span(), "';'"))?;
416 }
417
418 Ok(())
419}
420
421#[inline]
425fn get_constraints_from_terms<N: Int>(
426 comparator: Comparator,
427 terms: Vec<GeneralPBTerm<N>>,
428 coeff_sum: N,
429 degree: N,
430) -> (PBConstraintEnum, Option<PBConstraintEnum>)
431where
432 i64: TryFrom<N>,
433 PBConstraintEnum: From<GeneralPBConstraint<N>>,
434{
435 let leq_constraint = if comparator != Comparator::GreaterEqual {
437 let negated_terms = terms
438 .iter()
439 .map(|term| GeneralPBTerm::new(-term.coeff.clone(), term.lit))
440 .collect();
441
442 Some(constraint_from_terms_and_coeff_sum(
443 negated_terms,
444 -degree.clone(),
445 coeff_sum.clone(),
446 ))
447 } else {
448 None
449 };
450
451 match comparator {
452 Comparator::GreaterEqual => (
453 constraint_from_terms_and_coeff_sum(terms, degree, coeff_sum),
454 None,
455 ),
456 Comparator::LessEqual => (leq_constraint.unwrap(), None),
457 Comparator::Equal => (
458 constraint_from_terms_and_coeff_sum(terms, degree, coeff_sum),
459 leq_constraint,
460 ),
461 }
462}