type_unification.rs

use std::collections::VecDeque;
use std::fmt::{Display};

#[derive(Clone, Debug, PartialEq)]
pub enum InferredType {
    U8,
    U16,
    Record(Vec<(String, InferredType)>),
    AllOf(Vec<InferredType>),
}

fn main() {

    let record = InferredType::Record(vec![
        ("field1".to_string(), InferredType::AllOf(vec![
            InferredType::U16,
            InferredType::U16,
        ])),
        ("field2".to_string(), InferredType::Record(vec![
            ("field3".to_string(), InferredType::U8),
            ("field4".to_string(), InferredType::U16),
        ])),
    ]);

    let record2 = InferredType::Record(vec![
        ("field1".to_string(), InferredType::U8),
        ("field2".to_string(), InferredType::Record(vec![
            ("field3".to_string(), InferredType::U8),
            ("field4".to_string(), InferredType::U16),
        ])),
    ]);

    let all_of = InferredType::AllOf(vec![
       record, record2
    ]);

    let result = handle_logic(&all_of).unwrap();

    let result = flatten(result);

    dbg!(result);

    //dbg!(result);
}


fn handle_logic(typ: &InferredType) -> Result<IntermediateState, String> {
    let mut state = IntermediateState::from_inferred_type(typ);
    let mut state_vec = VecDeque::new();
    state_vec.push_back(&mut state);

    while let Some(state) = state_vec.pop_back() {
        match state {
            IntermediateState::Record {value} => {
                for (_, state) in value.iter_mut(){
                    let new_state = resolve_all_of(state)?;

                    *state = vec![new_state];

                    state_vec.extend(state);

                }
            }

             IntermediateState::Internal{ resolved_simple: resolved, unresolved} => {

                match resolved {
                    Some(resolved)  => {
                        // If resolved is Some, we can skip the rest
                       *unresolved = vec![];
                    }

                    None => {
                        let new_type = resolve_all_of(unresolved)?;
                        match new_type.get_resolved() {
                            Some(x) => {
                                *resolved = Some(x);
                                *unresolved = vec![];
                            },
                            None => {
                                *unresolved = vec![new_type];
                                state_vec.extend(unresolved);
                            }
                        }
                    }
                }

            }
        }
    }

    Ok(state)
}

fn resolve_all_of(types: &mut Vec<IntermediateState>) -> Result<IntermediateState, String> {
    let (first, rest) = types.split_at_mut(1);

    let mut state = first[0].clone();

    for typ in rest.iter_mut() {
        let mut result = handle_both(&mut state, typ)?;
        if let Some(new_state) = result {
            state = new_state;
        }
    }

    Ok(state)
}

#[derive(Debug, Clone)]
pub enum IntermediateState {
    Record {
        value: Vec<(String, Vec<IntermediateState>)>,
    },
    Internal {
        unresolved: Vec<IntermediateState>,
        resolved_simple: Option<InferredType>,
    }
}

impl IntermediateState {
    pub fn get_resolved(&self) -> Option<InferredType> {
        match self {
            IntermediateState::Internal { resolved_simple: resolved, .. }=> resolved.clone(),
            _ => None,
        }
    }

    pub fn from_inferred_type(typ: &InferredType) -> IntermediateState {
        match typ {
            InferredType::U8 => IntermediateState::Internal {
                unresolved: vec![],
                resolved_simple: Some(InferredType::U8),
            },
            InferredType::U16 => IntermediateState::Internal {
                unresolved: vec![],
                resolved_simple: Some(InferredType::U16),
            },
            InferredType::Record(fields) => IntermediateState::Record { value:
                fields.iter().map(|(field, inferred_type)| {
                    (field.clone(), vec![IntermediateState::from_inferred_type(inferred_type)])
                }).collect::<Vec<_>>()
            },
            InferredType::AllOf(res) => IntermediateState::Internal {
                resolved_simple: None,
                unresolved: res.iter().map( | typ| IntermediateState::from_inferred_type(typ)).collect:: < Vec<_ > > ()
            },
        }
    }

    pub fn get_record(&mut self) -> Option<&mut Vec<(String,Vec<IntermediateState>)>> {
        match self {
            IntermediateState::Record {value} => Some(value),
            _ => None,
        }
    }
}

// Returns None if updated in place, otherwise returns a new state
fn handle_both(left: &mut IntermediateState, right: &mut IntermediateState) -> Result<Option<IntermediateState>, String> {
    match (left, right) {
        (IntermediateState::Record { value: value1 }, IntermediateState::Record { value: value2 }) => {

            for (field_a, field_typ) in value1.iter_mut() {
                let result =
                    value2.iter().find(|(field_b, field_typ)| field_a == field_b);

                if let Some((_, inferred_type)) = result {
                    field_typ.extend(inferred_type.clone());
                } else {
                    return Err(format!("Field {} not found in right record", field_a));
                }
            }

            for (field_b, field_typ) in value2.iter() {
                let result =
                    value1.iter().find(|(field_a, field_typ)| field_a == field_b);

                if result.is_none() {
                    return Err(format!("Field {} not found in left record", field_b));
                }
            }

           Ok(None)
        }


        (IntermediateState::Internal { resolved_simple: Some(typ1), ..}, IntermediateState::Internal{ resolved_simple: Some(typ2), ..}) => {
            if typ1 == typ2 {
               Ok(None)
            } else {
                Err(format!("type mismatch: {:?} and {:?}", typ1, typ2))
            }
        }

        (IntermediateState::Internal { resolved_simple: None, unresolved: types}, IntermediateState::Internal{ resolved_simple: Nones, unresolved: types2}) => {

            types.extend(types2.clone());

            Ok(None)
        }

        (IntermediateState::Internal { resolved_simple: None, unresolved: types}, IntermediateState::Internal{ resolved_simple: Some(typ), ..}) => {
            types.push(IntermediateState::from_inferred_type(typ));
            Ok(None)
        }

        ( IntermediateState::Internal{ resolved_simple: Some(typ), ..}, IntermediateState::Internal { resolved_simple: None, unresolved: types}) => {
            // Modify `left` in place without dereferencing it
            types.push(IntermediateState::from_inferred_type(typ));
            Ok(Some(IntermediateState::Internal {
                unresolved: types.clone(),
                resolved_simple: None,
            }))

        }

        (intermediate_state1, intermediate_state2) => {
            Err(format!("Cannot handle types: {:?} and {:?}", intermediate_state1, intermediate_state2))
        }
    }
}


enum Task {
    Process {
        state: IntermediateState,
        path: Vec<String>,
    },
    BuildRecord {
        path: Vec<String>,
        fields: Vec<(String, InferredType)>,
    },
}



fn flatten(root: IntermediateState) -> Result<InferredType, String> {
    let mut queue = VecDeque::new();
    let mut results = std::collections::HashMap::new();

    queue.push_back(Task::Process {
        state: root,
        path: vec![],
    });

    while let Some(task) = queue.pop_back() {
        match task {
            Task::Process { state, path } => match state {
                IntermediateState::Internal { resolved_simple: Some(typ), .. } => {
                    results.insert(path, typ);
                }
                IntermediateState::Internal { resolved_simple: None, unresolved } => {
                    return Err(format!("Unresolved type: {:?}", unresolved));
                }
                IntermediateState::Record { value } => {
                    let mut subfields = Vec::new();

                    for (field_name, field_states) in value {
                        for field_state in field_states {
                            let mut new_path = path.clone();
                            new_path.push(field_name.clone());
                            queue.push_back(Task::Process {
                                state: field_state, // <-- now field_state is a single IntermediateState
                                path: new_path.clone(),
                            });
                            subfields.push((field_name.clone(), new_path));
                        }
                    }

                    queue.push_back(Task::BuildRecord {
                        path,
                        fields: subfields,
                    });
                }
            },
            Task::BuildRecord { path, fields } => {
                let mut collected = Vec::new();
                for (field_name, field_path) in fields {
                    let typ = results
                        .remove(&field_path)
                        .ok_or_else(|| format!("Missing field at {:?}", field_path))?;
                    collected.push((field_name, typ));
                }
                results.insert(path, InferredType::Record(collected));
            }
        }
    }

    results
        .remove(&Vec::new())
        .ok_or_else(|| "Final result missing".to_string())
}