codama-korok-visitors 0.9.2

Visitors for Korok trees
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122

use crate::KorokVisitor;
use codama_attributes::{Attributes, DefaultValueDirective, TryFromFilter};
use codama_errors::CodamaResult;
use codama_koroks::{KorokMut, KorokTrait};
use codama_nodes::{
    InstructionArgumentNode, Node, RegisteredTypeNode, StructFieldTypeNode, ValueNode,
};

#[derive(Default)]
pub struct SetDefaultValuesVisitor;

impl SetDefaultValuesVisitor {
    pub fn new() -> Self {
        Self
    }
}

impl KorokVisitor for SetDefaultValuesVisitor {
    fn visit_root(&mut self, korok: &mut codama_koroks::RootKorok) -> CodamaResult<()> {
        self.visit_children(korok)?;
        Ok(())
    }

    fn visit_crate(&mut self, korok: &mut codama_koroks::CrateKorok) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }

    fn visit_file_module(
        &mut self,
        korok: &mut codama_koroks::FileModuleKorok,
    ) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }

    fn visit_module(&mut self, korok: &mut codama_koroks::ModuleKorok) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }

    fn visit_struct(&mut self, korok: &mut codama_koroks::StructKorok) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }

    fn visit_enum(&mut self, korok: &mut codama_koroks::EnumKorok) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }

    fn visit_unsupported_item(
        &mut self,
        korok: &mut codama_koroks::UnsupportedItemKorok,
    ) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }

    fn visit_enum_variant(
        &mut self,
        korok: &mut codama_koroks::EnumVariantKorok,
    ) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }

    fn visit_field(&mut self, korok: &mut codama_koroks::FieldKorok) -> CodamaResult<()> {
        self.visit_children(korok)?;
        set_default_values(korok.into())
    }
}

fn set_default_values(mut korok: KorokMut) -> CodamaResult<()> {
    // Ensure it has attributes.
    let Some(attributes) = korok.attributes() else {
        return Ok(());
    };

    // Ensure there is a node to set a default value on.
    let Some(node) = get_node_with_default_value(korok.node(), attributes)? else {
        return Ok(());
    };

    korok.set_node(Some(node));
    Ok(())
}

fn get_node_with_default_value(
    node: &Option<Node>,
    attributes: &Attributes,
) -> CodamaResult<Option<Node>> {
    let Some(directive) = attributes.get_last(DefaultValueDirective::filter) else {
        return Ok(None);
    };
    let resolved_node = directive.node.try_resolved()?;

    match node {
        // Handle struct fields.
        Some(Node::Type(RegisteredTypeNode::StructField(field))) => {
            let value = ValueNode::try_from(resolved_node.clone()).ok();
            Ok(value.map(|value| {
                StructFieldTypeNode {
                    default_value: Some(value),
                    default_value_strategy: directive.default_value_strategy,
                    ..field.clone()
                }
                .into()
            }))
        }
        // Handle instruction arguments.
        Some(Node::InstructionArgument(argument)) => Ok(Some(
            InstructionArgumentNode {
                default_value: Some(resolved_node.clone()),
                default_value_strategy: directive.default_value_strategy,
                ..argument.clone()
            }
            .into(),
        )),
        _ => Ok(None),
    }
}