shape/case_enum/

all.rs

use crate::case_enum::deduplicate_shape::extend_or_insert_shape;
use crate::location::Location;
use crate::{Shape, ShapeCase};
use indexmap::{IndexMap, IndexSet};

/// The inner type of [`ShapeCase::All`] which intentionally can't be constructed directly because
/// it requires simplifying the inner shapes to achieve the desired result.
///
/// Construct this via [`Shape::all`].
#[derive(Clone, Debug, Eq)]
pub struct All {
    shapes: Vec<Shape>,
}

impl All {
    pub fn iter(&self) -> impl Iterator<Item = &Shape> {
        self.shapes.iter()
    }
}

impl<'a> IntoIterator for &'a All {
    type Item = &'a Shape;
    type IntoIter = std::slice::Iter<'a, Shape>;

    fn into_iter(self) -> Self::IntoIter {
        self.shapes.iter()
    }
}

impl PartialEq<All> for All {
    fn eq(&self, other: &All) -> bool {
        self.shapes.iter().collect::<IndexSet<_>>() == other.shapes.iter().collect::<IndexSet<_>>()
    }
}

/// Builds a [`Shape`] with [`ShapeCase::All`] _unless_ the shapes can simplify to a single
/// [`ShapeCase`].
#[allow(clippy::too_many_lines)]
pub(crate) fn all(shapes: impl Iterator<Item = Shape>, mut locations: Vec<Location>) -> Shape {
    let mut flat_shapes = Vec::new();
    for shape in shapes {
        if let ShapeCase::All(inner_shapes) = shape.case() {
            locations.extend(shape.locations.clone());
            for inner_shape in inner_shapes {
                extend_or_insert_shape(&mut flat_shapes, inner_shape.clone());
            }
        } else {
            extend_or_insert_shape(&mut flat_shapes, shape);
        }
    }

    // Store arrays and objects separately so we can merge them
    let mut arrays_to_merge = Vec::new();
    let mut objects_to_merge = Vec::new();
    flat_shapes.retain(|shape| match shape.case() {
        ShapeCase::Array { prefix, tail } => {
            // Store arrays separately so we can merge them into one array
            arrays_to_merge.push((prefix.clone(), tail.clone(), shape.locations.clone()));
            false
        }
        ShapeCase::Object { fields, rest } => {
            // Store object fields separately so we can merge them into one object
            objects_to_merge.push((fields.clone(), rest.clone(), shape.locations.clone()));
            false
        }
        _ => true,
    });

    // The presence of the null literal shape (ShapeCase::Null)
    // in flat_shapes subsumes all the other shapes.
    if let Some(null_shape) = flat_shapes
        .iter()
        .find(|shape| shape.case() == ShapeCase::Null)
    {
        return null_shape.clone();
    }

    // The presence of any specific Bool literal shapes subsumes
    // the general Bool shape.
    if flat_shapes
        .iter()
        .any(|shape| matches!(shape.case.as_ref(), ShapeCase::Bool(Some(_))))
    {
        flat_shapes.retain(|shape| !matches!(shape.case.as_ref(), ShapeCase::Bool(None)));
    }

    if flat_shapes
        .iter()
        .any(|shape| matches!(shape.case.as_ref(), ShapeCase::Int(Some(_))))
    {
        // The presence of any specific Int literal shapes subsumes
        // the general Int shape and the general Float shape.
        flat_shapes.retain(|shape| {
            !matches!(shape.case.as_ref(), ShapeCase::Float | ShapeCase::Int(None))
        });
    } else if flat_shapes
        .iter()
        .any(|shape| shape.case() == ShapeCase::Int(None))
    {
        // The presence of the general Int shape subsumes the general
        // Float shape.
        flat_shapes.retain(|shape| !matches!(shape.case.as_ref(), ShapeCase::Float));
    }

    // The presence of any specific String literal shapes
    // subsumes the general String shape.
    if flat_shapes
        .iter()
        .any(|shape| matches!(shape.case.as_ref(), ShapeCase::String(Some(_))))
    {
        flat_shapes.retain(|shape| !matches!(shape.case.as_ref(), ShapeCase::String(None)));
    }

    // Merge arrays and insert into flat_shapes
    if !arrays_to_merge.is_empty() {
        let longest_prefix_len = arrays_to_merge
            .iter()
            .fold(0, |acc, (items, _, _)| acc.max(items.len()));

        let mut array_locations = Vec::new();

        // First, merge the ShapeCase::Array prefix vectors,
        // combining overlapping elements using ShapeCase::All.
        let fixed_element_shapes = (0..longest_prefix_len)
            .map(|i| {
                let mut element_shapes = Vec::new();
                let mut element_locations = Vec::new();
                for (items, rest, locations) in &arrays_to_merge {
                    array_locations.extend(locations.clone());
                    element_shapes.push(items.get(i).unwrap_or(rest).clone());
                    element_locations.extend_from_slice(&items.get(i).unwrap_or(rest).locations);
                }
                Shape::all(element_shapes, element_locations)
            })
            .collect::<Vec<_>>();

        // Second, merge the ShapeCase::Array rest vectors,
        // combining non-None shapes using ShapeCase::All.
        let mut rest_shapes_to_merge = Vec::new();
        for rest in arrays_to_merge
            .iter()
            .filter_map(|(_items, rest, _locations)| match rest.case.as_ref() {
                ShapeCase::None => None,
                _ => Some(rest.clone()),
            })
        {
            extend_or_insert_shape(&mut rest_shapes_to_merge, rest);
        }

        let merged_rest = if rest_shapes_to_merge.is_empty() {
            Shape::none()
        } else {
            // There are no locations carried forward here because this is a virtual All, not
            // one created directly by user code.
            all(rest_shapes_to_merge.into_iter(), Vec::new())
        };

        // Finally, insert the merged ShapeCase::Array into
        // flat_shapes.
        extend_or_insert_shape(
            &mut flat_shapes,
            Shape::array(fixed_element_shapes, merged_rest, array_locations),
        );
    }

    // Merge object fields and insert into flat_shapes
    if !objects_to_merge.is_empty() {
        let mut merged_object_fields = IndexMap::new();
        let mut merged_object_rest = Shape::none();
        let mut object_locations = Vec::new();

        for (fields, rest, locations) in objects_to_merge {
            for (field_name, field_shape) in fields {
                merged_object_fields
                    .entry(field_name)
                    .or_insert_with(Vec::new)
                    .push(field_shape);
            }
            object_locations.extend(locations);

            merged_object_rest = match (merged_object_rest.case(), rest.case()) {
                (ShapeCase::None, _) => rest.clone(),
                (_, ShapeCase::None) => merged_object_rest.clone(),
                _ => all([merged_object_rest, rest.clone()].into_iter(), Vec::new()),
            };
        }

        flat_shapes.push(Shape::object(
            merged_object_fields
                .into_iter()
                .map(|(k, v)| (k, all(v.into_iter(), Vec::new())))
                .collect(),
            merged_object_rest,
            object_locations,
        ));
    }

    // An empty ::All intersection is equivalent to Unknown, since
    // it imposes no expectations or criteria on input shapes.
    if flat_shapes.is_empty() {
        return Shape::unknown(locations);
    }

    // Remove None from flat_shapes, since it has no marginal effect
    // on ::All intersections.
    flat_shapes.retain(|shape| !shape.is_none());
    // If removing None resulted in empty flat_shapes, return
    // None to represent an empty intersection.
    if flat_shapes.is_empty() {
        return Shape::none();
    }

    // Similar logic but for Unknown, which also (like None)
    // contributes nothing to the intersection.
    flat_shapes.retain(|shape| !shape.is_unknown());
    if flat_shapes.is_empty() {
        return Shape::unknown(locations);
    }

    // If there's only one shape in the intersection, return it.
    if let (Some(only_shape), 1) = (flat_shapes.first(), flat_shapes.len()) {
        return only_shape.clone();
    }

    // To bias simplification towards disjunctive normal form
    // (unions of intersections), unpack any ShapeCase::One items
    // contained by flat_shapes.
    //
    // Unpacking the ShapeCase::One items results in a
    // ShapeCase::One union of ShapeCase::All shapes (i.e.
    // disjunctive normal form) where each ShapeCase::All shape
    // represents a distinct choice from each available union,
    // treating single non-union shapes as unions of one.
    //
    // For example
    //
    //   (A | B | C) & D & (E | F) =>
    //     (A & D & E) | (A & D & F) |
    //     (B & D & E) | (B & D & F) |
    //     (C & D & E) | (C & D & F)
    //
    // Although D is not a ShapeCase::One union, it can be treated
    // as a union of just one member shape (D) for the purposes of
    // this enumeration.
    let new_all_shapes = flat_shapes
        .iter()
        .fold(vec![vec![]], |acc, shape| {
            let mut new_shapes = Vec::new();
            for partial_prefix in &acc {
                if let ShapeCase::One(one_shapes) = shape.case.as_ref() {
                    for one_shape in one_shapes {
                        let mut new_partial = partial_prefix.clone();
                        if !one_shape.is_none() {
                            new_partial.push(one_shape.clone());
                        }
                        new_shapes.push(new_partial);
                    }
                } else {
                    let mut new_partial = partial_prefix.clone();
                    if !shape.is_none() {
                        new_partial.push(shape.clone());
                    }
                    new_shapes.push(new_partial);
                }
            }
            new_shapes
        })
        .into_iter()
        .collect::<IndexSet<_>>();

    if new_all_shapes.len() == 1 {
        Shape::new(
            ShapeCase::All(All {
                shapes: flat_shapes,
            }),
            locations,
        )
    } else {
        let simplified_all_shapes = new_all_shapes
            .into_iter()
            .map(|shapes| Shape::all(shapes, []))
            .collect::<Vec<Shape>>();
        Shape::one(simplified_all_shapes, locations)
    }
}