tree_sitter_stack_graphs/
functions.rs

1// -*- coding: utf-8 -*-
2// ------------------------------------------------------------------------------------------------
3// Copyright © 2021, stack-graphs authors.
4// Licensed under either of Apache License, Version 2.0, or MIT license, at your option.
5// Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
6// ------------------------------------------------------------------------------------------------
7
8//! Define tree-sitter-graph functions
9
10pub use path::add_path_functions;
11
12pub mod path {
13    use std::path::Component;
14    use std::path::Path;
15    use std::path::PathBuf;
16    use tree_sitter_graph::functions::Function;
17    use tree_sitter_graph::functions::Functions;
18    use tree_sitter_graph::functions::Parameters;
19    use tree_sitter_graph::graph::Graph;
20    use tree_sitter_graph::graph::Value;
21    use tree_sitter_graph::ExecutionError;
22
23    pub fn add_path_functions(functions: &mut Functions) {
24        functions.add(
25            "path-dir".into(),
26            path_fn(|p| p.parent().map(|s| s.as_os_str().to_os_string())),
27        );
28        functions.add(
29            "path-fileext".into(),
30            path_fn(|p| p.extension().map(|s| s.to_os_string())),
31        );
32        functions.add(
33            "path-filename".into(),
34            path_fn(|p| p.file_name().map(|s| s.to_os_string())),
35        );
36        functions.add(
37            "path-filestem".into(),
38            path_fn(|p| p.file_stem().map(|s| s.to_os_string())),
39        );
40        functions.add("path-join".into(), PathJoin);
41        functions.add(
42            "path-normalize".into(),
43            path_fn(|p| normalize(p).map(|p| p.as_os_str().to_os_string())),
44        );
45        functions.add("path-split".into(), PathSplit);
46    }
47
48    pub fn path_fn<F>(f: F) -> impl Function
49    where
50        F: Fn(&Path) -> Option<std::ffi::OsString>,
51    {
52        PathFn(f)
53    }
54
55    struct PathFn<F>(F)
56    where
57        F: Fn(&Path) -> Option<std::ffi::OsString>;
58
59    impl<F> Function for PathFn<F>
60    where
61        F: Fn(&Path) -> Option<std::ffi::OsString>,
62    {
63        fn call(
64            &self,
65            _graph: &mut Graph,
66            _source: &str,
67            parameters: &mut dyn Parameters,
68        ) -> Result<Value, ExecutionError> {
69            let path = PathBuf::from(parameters.param()?.into_string()?);
70            parameters.finish()?;
71
72            let path = self.0(&path);
73            Ok(path
74                .map(|s| {
75                    s.into_string()
76                        .unwrap_or_else(|s| s.to_string_lossy().to_string())
77                        .into()
78                })
79                .unwrap_or(Value::Null))
80        }
81    }
82
83    struct PathJoin;
84
85    impl Function for PathJoin {
86        fn call(
87            &self,
88            _graph: &mut Graph,
89            _source: &str,
90            parameters: &mut dyn Parameters,
91        ) -> Result<Value, ExecutionError> {
92            let mut path = PathBuf::new();
93            while let Ok(component) = parameters.param() {
94                path = path.join(component.into_string()?);
95            }
96
97            Ok(path.to_str().unwrap().into())
98        }
99    }
100
101    struct PathSplit;
102
103    impl Function for PathSplit {
104        fn call(
105            &self,
106            _graph: &mut Graph,
107            _source: &str,
108            parameters: &mut dyn Parameters,
109        ) -> Result<Value, ExecutionError> {
110            let path = PathBuf::from(parameters.param()?.into_string()?);
111            parameters.finish()?;
112
113            let components = path
114                .components()
115                .map(|c| c.as_os_str().to_str().unwrap().into())
116                .collect::<Vec<_>>();
117            Ok(components.into())
118        }
119    }
120
121    /// Normalize a path, removing things like `.` and `..` wherever possible.
122    // Based on the following code from Cargo:
123    // https://github.com/rust-lang/cargo/blob/e515c3277bf0681bfc79a9e763861bfe26bb05db/crates/cargo-util/src/paths.rs#L73-L106
124    // Licensed under MIT license & Apache License (Version 2.0)
125    pub fn normalize(path: &Path) -> Option<PathBuf> {
126        let mut components = path.components().peekable();
127        let mut ret = if let Some(c @ Component::Prefix(..)) = components.peek().cloned() {
128            components.next();
129            PathBuf::from(c.as_os_str())
130        } else {
131            PathBuf::new()
132        };
133
134        let mut has_root = false;
135        let mut normal_components = 0usize;
136        for component in components {
137            match component {
138                Component::Prefix(..) => unreachable!(),
139                Component::RootDir => {
140                    has_root = true;
141                    ret.push(component.as_os_str());
142                }
143                Component::CurDir => {}
144                Component::ParentDir => {
145                    if normal_components > 0 {
146                        normal_components -= 1;
147                        ret.pop();
148                    } else if has_root {
149                        return None;
150                    } else {
151                        ret.push(component.as_os_str());
152                    }
153                }
154                Component::Normal(c) => {
155                    normal_components += 1;
156                    ret.push(c);
157                }
158            }
159        }
160        Some(ret)
161    }
162}