sigmd 0.1.0

//! Lowering of decoded SAL annotations into model buffers.
//!
//! Each recognized annotation has a [`Rule`] that names it and the
//! buffer slots it emits. Adding a new annotation is one entry in
//! [`RULES`].

use std::collections::HashMap;

use sigmd::model::{
    BinaryExpression, BinaryOperator, Buffer, BufferDirection, BufferPhase, Expression, Parameter,
    UnaryExpression, UnaryOperator,
};

use super::{annotation::Annotation, parse};
use crate::cli::compiler::BuildContext;

/// Runs analysis over a function's parameters and returns the derived
/// buffer descriptors.
pub fn analyze(
    parameters: &[Parameter],
    annotations_per_parameter: &[Vec<Annotation<'_>>],
    ctx: &BuildContext,
) -> Vec<Buffer> {
    let resolver = Resolver::new(parameters, &ctx.sizeofs);

    // Lower every SAL annotation to a raw buffer descriptor. One parameter
    // may produce several entries, since clang merges redeclarations onto
    // a single `ParmVarDecl` and `inject.h` strips `_When_(Level == N, ...)`
    // from tagged-union APIs.
    let mut produced = Vec::new();
    for (index, annotations) in annotations_per_parameter.iter().enumerate() {
        assert!(index <= u8::MAX as usize);
        let parameter_index = index as u8;

        for annotation in annotations {
            produced.extend(buffers_for(parameter_index, annotation, &resolver));
        }
    }

    // Group by `(parameter, direction, phase)` so each group can be
    // reconciled independently in the pass below.
    let mut groups = Vec::<(BufferKey, Vec<Buffer>)>::new();
    for buffer in produced {
        let key = BufferKey::from(&buffer);
        match groups.iter_mut().find(|(other, _)| *other == key) {
            Some((_, group)) => group.push(buffer),
            None => groups.push((key, vec![buffer])),
        }
    }

    // Collapse groups whose members agree on `length`. Drop the rest.
    //
    // Picking one variant arbitrarily would mislead the consumers, since
    // the right answer depends on a runtime sibling argument the metadata
    // does not model.
    let mut output = Vec::new();
    for (key, group) in groups {
        let count = group.len();
        let mut iter = group.into_iter();
        let first = iter.next().expect("group is non-empty by construction");
        if iter.all(|buffer| buffer.length == first.length) {
            output.push(first);
        }
        else {
            tracing::debug!(
                parameter = key.parameter,
                direction = ?key.direction,
                phase = ?key.phase,
                variants = count,
                "dropping ambiguous buffer group with conflicting lengths"
            );
        }
    }

    output
}

/// Identity used to group buffer descriptors before reconciliation.
/// Two buffers with equal keys target the same argument slot at the
/// same call phase, so the analyzer must agree on a single `length`
/// for them or drop the group.
#[derive(Clone, Copy, PartialEq, Eq)]
struct BufferKey {
    /// Index of the parameter the buffer attaches to.
    parameter: u8,

    /// Whether the buffer is read or written by the call.
    direction: BufferDirection,

    /// Pre- or post-call evaluation phase.
    phase: BufferPhase,
}

impl From<&Buffer> for BufferKey {
    fn from(buffer: &Buffer) -> Self {
        Self {
            parameter: buffer.parameter,
            direction: buffer.direction,
            phase: buffer.phase,
        }
    }
}

/// Per-function context for lowering parser expressions to model expressions.
struct Resolver<'a> {
    /// Parameter index by name. Built once per function.
    parameter_by_name: HashMap<&'a str, u8>,

    /// Type-name to byte-size lookup.
    sizeofs: &'a HashMap<String, u64>,
}

impl<'a> Resolver<'a> {
    /// Builds a resolver over `parameters` and the per-TU sizeof table.
    fn new(parameters: &'a [Parameter], sizeofs: &'a HashMap<String, u64>) -> Self {
        let mut parameter_by_name = HashMap::with_capacity(parameters.len());
        for (index, parameter) in parameters.iter().enumerate() {
            assert!(index <= u8::MAX as usize);
            if let Some(name) = parameter.name.as_deref() {
                parameter_by_name.insert(name, index as u8);
            }
        }

        Self {
            parameter_by_name,
            sizeofs,
        }
    }

    /// Resolves a parameter name to its index.
    fn parameter(&self, name: &str) -> Option<u8> {
        let index = self.parameter_by_name.get(name).copied();
        if index.is_none() {
            tracing::debug!(parameter = name, "SAL referenced unknown parameter");
        }
        index
    }

    /// Resolves a type name to its byte size.
    fn sizeof(&self, name: &str) -> Option<u64> {
        let size = self.sizeofs.get(name).copied();
        if size.is_none() {
            tracing::debug!(ty = name, "sizeof() of unknown type");
        }
        size
    }
}

/// One emitted buffer slot for a recognized annotation.
struct Emit {
    /// Index into [`Annotation::args`] to lower as the buffer length.
    arg: usize,

    /// Buffer direction this emit produces.
    direction: BufferDirection,

    /// Pre- or post-call phase for the length expression.
    phase: BufferPhase,
}

/// Mapping from annotation name to its buffer-emit shape.
struct Rule {
    /// Annotation name as it appears in [`Annotation::name`].
    name: &'static str,

    /// One [`Emit`] per buffer this annotation should produce.
    emits: &'static [Emit],
}

/// Recognized SAL annotations and the buffer slots they emit.
///
//
// Buffer annotations
// The `size` argument is valid before the function call.
// The `count` argument is valid after the function call.
//
const RULES: &[Rule] = &[
    // _In_reads_bytes_(size)
    Rule {
        name: "_In_reads_bytes_",
        emits: &[Emit {
            arg: 0,
            direction: BufferDirection::Input,
            phase: BufferPhase::Pre,
        }],
    },
    // _In_reads_bytes_opt_(size)
    Rule {
        name: "_In_reads_bytes_opt_",
        emits: &[Emit {
            arg: 0,
            direction: BufferDirection::Input,
            phase: BufferPhase::Pre,
        }],
    },
    // _Out_writes_bytes_(size)
    Rule {
        name: "_Out_writes_bytes_",
        emits: &[Emit {
            arg: 0,
            direction: BufferDirection::Output,
            phase: BufferPhase::Pre,
        }],
    },
    // _Out_writes_bytes_opt_(size)
    Rule {
        name: "_Out_writes_bytes_opt_",
        emits: &[Emit {
            arg: 0,
            direction: BufferDirection::Output,
            phase: BufferPhase::Pre,
        }],
    },
    // _Inout_updates_bytes_(size)
    // intentional: `_Inout_updates_bytes_` (no `_to_`) emits only the
    //              output side. The input-read aspect is dropped.
    Rule {
        name: "_Inout_updates_bytes_",
        emits: &[Emit {
            arg: 0,
            direction: BufferDirection::Output,
            phase: BufferPhase::Pre,
        }],
    },
    // _Inout_updates_bytes_opt_(size)
    // intentional: `_Inout_updates_bytes_opt_` (no `_to_`) emits only the
    //              output side. The input-read aspect is dropped.
    Rule {
        name: "_Inout_updates_bytes_opt_",
        emits: &[Emit {
            arg: 0,
            direction: BufferDirection::Output,
            phase: BufferPhase::Pre,
        }],
    },
    // _Out_writes_bytes_to_(size, count)
    Rule {
        name: "_Out_writes_bytes_to_",
        emits: &[Emit {
            arg: 1,
            direction: BufferDirection::Output,
            phase: BufferPhase::Post,
        }],
    },
    // _Out_writes_bytes_to_opt_(size, count)
    Rule {
        name: "_Out_writes_bytes_to_opt_",
        emits: &[Emit {
            arg: 1,
            direction: BufferDirection::Output,
            phase: BufferPhase::Post,
        }],
    },
    // _Inout_updates_bytes_to_(size, count)
    Rule {
        name: "_Inout_updates_bytes_to_",
        emits: &[
            Emit {
                arg: 0,
                direction: BufferDirection::Input,
                phase: BufferPhase::Pre,
            },
            Emit {
                arg: 1,
                direction: BufferDirection::Output,
                phase: BufferPhase::Post,
            },
        ],
    },
    // _Inout_updates_bytes_to_opt_(size, count)
    Rule {
        name: "_Inout_updates_bytes_to_opt_",
        emits: &[
            Emit {
                arg: 0,
                direction: BufferDirection::Input,
                phase: BufferPhase::Pre,
            },
            Emit {
                arg: 1,
                direction: BufferDirection::Output,
                phase: BufferPhase::Post,
            },
        ],
    },
    // `_COM_Outptr_*` annotations are recognized via
    // `ParameterFlags::HAS_COM_ATTRIBUTE` in `super::annotation`. They do
    // not emit a buffer.
    //
    // Rule { name: "_COM_Outptr_",                       emits: &[] },
    // Rule { name: "_COM_Outptr_opt_",                   emits: &[] },
    // Rule { name: "_COM_Outptr_result_maybenull_",      emits: &[] },
    // Rule { name: "_COM_Outptr_opt_result_maybenull_",  emits: &[] },
];

/// Applies the matching rule for an annotation and returns the buffers
/// it emits.
fn buffers_for(
    parameter_index: u8,
    annotation: &Annotation<'_>,
    resolver: &Resolver<'_>,
) -> Vec<Buffer> {
    let rule = match RULES.iter().find(|rule| rule.name == annotation.name) {
        Some(rule) => rule,
        None => return Vec::new(),
    };

    rule.emits
        .iter()
        .filter_map(|emit| build_buffer(parameter_index, annotation, emit, resolver))
        .collect()
}

/// Lowers a single emit to a [`Buffer`], returning `None` on parse or
/// resolution failure.
fn build_buffer(
    parameter_index: u8,
    annotation: &Annotation<'_>,
    emit: &Emit,
    resolver: &Resolver<'_>,
) -> Option<Buffer> {
    let arg = match annotation.args.get(emit.arg) {
        Some(arg) => arg,
        None => {
            tracing::debug!(
                annotation = annotation.name,
                arg = emit.arg,
                "missing SAL argument"
            );
            return None;
        }
    };

    let parsed = match parse::parse(arg) {
        Ok(parsed) => parsed,
        Err(err) => {
            tracing::debug!(
                %err,
                annotation = annotation.name,
                "SAL arg parse failed"
            );
            return None;
        }
    };

    let length = lower(parsed, resolver)?;

    Some(
        Buffer::builder()
            .parameter(parameter_index)
            .length(length)
            .direction(emit.direction)
            .phase(emit.phase)
            .build(),
    )
}

/// Lowers a parser expression to a model expression.
fn lower(expr: parse::Expression<'_>, resolver: &Resolver<'_>) -> Option<Expression> {
    match expr {
        parse::Expression::Return => Some(Expression::Return),
        parse::Expression::Constant(value) => Some(Expression::Constant(value)),
        parse::Expression::Identifier(name) => {
            Some(Expression::Parameter(resolver.parameter(name)?))
        }
        parse::Expression::UnaryExpression(unary) => lower_unary(unary, resolver),
        parse::Expression::BinaryExpression(binary) => lower_binary(binary, resolver),
    }
}

/// Lowers a unary expression. `sizeof(IDENT)` resolves to a constant via
/// the resolver. Any other operand under `sizeof(...)` is dropped with
/// a trace, since the resolver only indexes by bare type names.
fn lower_unary(unary: parse::UnaryExpression<'_>, resolver: &Resolver<'_>) -> Option<Expression> {
    match unary.operator {
        parse::UnaryOperator::SizeOf => {
            let name = match *unary.expression {
                parse::Expression::Identifier(name) => name,
                _ => {
                    tracing::debug!("sizeof() requires a bare identifier");
                    return None;
                }
            };
            Some(Expression::Constant(resolver.sizeof(name)?))
        }
        parse::UnaryOperator::Dereference => {
            let inner = lower(*unary.expression, resolver)?;
            Some(Expression::UnaryExpression(Box::new(
                UnaryExpression::builder()
                    .operator(UnaryOperator::Dereference)
                    .expression(inner)
                    .build(),
            )))
        }
    }
}

/// Lowers a binary expression, propagating any leaf failure.
fn lower_binary(
    binary: parse::BinaryExpression<'_>,
    resolver: &Resolver<'_>,
) -> Option<Expression> {
    let lhs = lower(*binary.lhs, resolver)?;
    let rhs = lower(*binary.rhs, resolver)?;
    let operator = match binary.operator {
        parse::BinaryOperator::Add => BinaryOperator::Add,
        parse::BinaryOperator::Subtract => BinaryOperator::Subtract,
        parse::BinaryOperator::Multiply => BinaryOperator::Multiply,
        parse::BinaryOperator::Divide => BinaryOperator::Divide,
    };
    Some(Expression::BinaryExpression(Box::new(
        BinaryExpression::builder()
            .operator(operator)
            .lhs(lhs)
            .rhs(rhs)
            .build(),
    )))
}

#[cfg(test)]
mod tests {
    use sigmd::model::{Type, TypeKind};

    use super::*;

    fn param(name: &str, kind: TypeKind) -> Parameter {
        Parameter::builder()
            .name(name)
            .ty(Type::builder().name(name).kind(kind).build())
            .build()
    }

    #[test]
    fn out_writes_bytes_to_emits_post_output_buffer() {
        let parameters = vec![
            param("hFile", TypeKind::U64),
            param("lpBuffer", TypeKind::U8),
            param("nNumberOfBytesToRead", TypeKind::U32),
            param("lpNumberOfBytesRead", TypeKind::U32),
        ];
        let annotations = vec![
            vec![],
            vec![Annotation {
                name: "_Out_writes_bytes_to_",
                args: vec!["nNumberOfBytesToRead", "*lpNumberOfBytesRead"],
            }],
            vec![],
            vec![],
        ];
        let buffers = analyze(&parameters, &annotations, &BuildContext::default());
        assert_eq!(buffers.len(), 1);
        assert_eq!(buffers[0].parameter, 1);
        assert_eq!(buffers[0].direction, BufferDirection::Output);
        assert_eq!(buffers[0].phase, BufferPhase::Post);
        match &buffers[0].length {
            Expression::UnaryExpression(unary) => {
                assert!(matches!(unary.operator, UnaryOperator::Dereference));
                assert!(matches!(unary.expression, Expression::Parameter(3)));
            }
            other => panic!("unexpected length expression: {other:?}"),
        }
    }

    #[test]
    fn inout_updates_bytes_to_emits_two_buffers() {
        let parameters = vec![
            param("buf", TypeKind::U8),
            param("nIn", TypeKind::U32),
            param("pnOut", TypeKind::U32),
        ];
        let annotations = vec![
            vec![Annotation {
                name: "_Inout_updates_bytes_to_",
                args: vec!["nIn", "*pnOut"],
            }],
            vec![],
            vec![],
        ];
        let buffers = analyze(&parameters, &annotations, &BuildContext::default());
        assert_eq!(buffers.len(), 2);
        assert_eq!(buffers[0].direction, BufferDirection::Input);
        assert_eq!(buffers[0].phase, BufferPhase::Pre);
        assert_eq!(buffers[1].direction, BufferDirection::Output);
        assert_eq!(buffers[1].phase, BufferPhase::Post);
    }

    #[test]
    fn sizeof_resolves_to_constant() {
        let parameters = vec![param("buf", TypeKind::U8)];
        let annotations = vec![vec![Annotation {
            name: "_In_reads_bytes_",
            args: vec!["sizeof(WIN32_FIND_DATAA)"],
        }]];

        let buffers = analyze(
            &parameters,
            &annotations,
            &BuildContext {
                sizeofs: HashMap::from([(String::from("WIN32_FIND_DATAA"), 0x0140)]),
                ..Default::default()
            },
        );
        assert_eq!(buffers.len(), 1);
        assert!(matches!(buffers[0].length, Expression::Constant(0x0140)));
    }

    #[test]
    fn unrecognized_annotation_emits_nothing() {
        let parameters = vec![param("p", TypeKind::U32)];
        let annotations = vec![vec![Annotation {
            name: "_Reserved_",
            args: vec![],
        }]];
        let buffers = analyze(&parameters, &annotations, &BuildContext::default());
        assert!(buffers.is_empty());
    }

    #[test]
    fn fully_identical_buffers_from_redeclaration_collapse_to_one() {
        // Two SAL annotations producing identical buffers. Models clang
        // merging two function redeclarations whose inject-header
        // per-call counters let both annotations through to the
        // analyzer with the same length expression.
        let parameters = vec![param("name", TypeKind::U8), param("namelen", TypeKind::U32)];
        let annotations = vec![
            vec![
                Annotation {
                    name: "_Out_writes_bytes_",
                    args: vec!["namelen"],
                },
                Annotation {
                    name: "_Out_writes_bytes_",
                    args: vec!["namelen"],
                },
            ],
            vec![],
        ];
        let buffers = analyze(&parameters, &annotations, &BuildContext::default());
        assert_eq!(buffers.len(), 1);
        assert_eq!(buffers[0].parameter, 0);
        assert_eq!(buffers[0].direction, BufferDirection::Output);
        assert_eq!(buffers[0].phase, BufferPhase::Pre);
    }

    #[test]
    fn ambiguous_buffers_with_conflicting_lengths_are_dropped() {
        // Tagged-union Win32 APIs like AddPrinterDriverEx attach multiple
        // _When_(Level == N, _In_reads_bytes_(sizeof(STRUCTN))) annotations
        // to one parameter. After inject.h strips _When_, the analyzer
        // sees several _In_reads_bytes_ with different constant lengths
        // sharing one (parameter, direction, phase). The right answer
        // depends on a sibling Level argument that the metadata does not
        // model, so any single pick would mislead the monitor. Drop the
        // entire group instead.
        let parameters = vec![param("buf", TypeKind::U8)];
        let annotations = vec![vec![
            Annotation {
                name: "_In_reads_bytes_",
                args: vec!["sizeof(SMALL)"],
            },
            Annotation {
                name: "_In_reads_bytes_",
                args: vec!["sizeof(BIG)"],
            },
        ]];
        let buffers = analyze(
            &parameters,
            &annotations,
            &BuildContext {
                sizeofs: HashMap::from([(String::from("SMALL"), 32), (String::from("BIG"), 128)]),
                ..Default::default()
            },
        );
        assert!(buffers.is_empty(), "expected drop, got {buffers:?}");
    }

    #[test]
    fn unknown_identifier_drops_the_buffer() {
        let parameters = vec![param("p", TypeKind::U32)];
        let annotations = vec![vec![Annotation {
            name: "_In_reads_bytes_",
            args: vec!["nGhost"],
        }]];
        let buffers = analyze(&parameters, &annotations, &BuildContext::default());
        assert!(buffers.is_empty());
    }
}