nixl_sys/descriptors/
query.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8// http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use super::*;
17use crate::Params;
18
19/// A safe wrapper around a NIXL query response list
20pub struct QueryResponseList {
21    inner: NonNull<bindings::nixl_capi_query_resp_list_s>,
22}
23
24/// Represents a single query response which may or may not contain parameters
25pub struct QueryResponse<'a> {
26    list: &'a QueryResponseList,
27    index: usize,
28}
29
30impl QueryResponseList {
31    /// Creates a new empty query response list
32    pub fn new() -> Result<Self, NixlError> {
33        let mut list = ptr::null_mut();
34        let status = unsafe { nixl_capi_create_query_resp_list(&mut list) };
35
36        match status {
37            NIXL_CAPI_SUCCESS => {
38                // SAFETY: If status is NIXL_CAPI_SUCCESS, list is non-null
39                let inner = unsafe { NonNull::new_unchecked(list) };
40                Ok(Self { inner })
41            }
42            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
43            _ => Err(NixlError::BackendError),
44        }
45    }
46
47    /// Returns the number of responses in the list
48    pub fn len(&self) -> Result<usize, NixlError> {
49        let mut size = 0;
50        let status = unsafe { nixl_capi_query_resp_list_size(self.inner.as_ptr(), &mut size) };
51
52        match status {
53            NIXL_CAPI_SUCCESS => Ok(size),
54            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
55            _ => Err(NixlError::BackendError),
56        }
57    }
58
59    /// Returns true if the list is empty
60    pub fn is_empty(&self) -> Result<bool, NixlError> {
61        Ok(self.len()? == 0)
62    }
63
64    /// Gets a query response at the given index
65    pub fn get(&self, index: usize) -> Result<QueryResponse<'_>, NixlError> {
66        let size = self.len()?;
67        if index >= size {
68            return Err(NixlError::InvalidParam);
69        }
70
71        Ok(QueryResponse { list: self, index })
72    }
73
74    /// Returns an iterator
75    pub fn iter(&self) -> Result<QueryResponseIterator<'_>, NixlError> {
76        Ok(QueryResponseIterator {
77            list: self,
78            index: 0,
79            len: self.len()?,
80        })
81    }
82
83    pub(crate) fn handle(&self) -> *mut bindings::nixl_capi_query_resp_list_s {
84        self.inner.as_ptr()
85    }
86}
87
88impl<'a> QueryResponse<'a> {
89    /// Returns true if this response contains parameters
90    pub fn has_value(&self) -> Result<bool, NixlError> {
91        let mut has_value = false;
92        let status = unsafe {
93            nixl_capi_query_resp_list_has_value(
94                self.list.inner.as_ptr(),
95                self.index,
96                &mut has_value,
97            )
98        };
99
100        match status {
101            NIXL_CAPI_SUCCESS => Ok(has_value),
102            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
103            _ => Err(NixlError::BackendError),
104        }
105    }
106
107    /// Gets the parameters if this response has a value
108    pub fn get_params(&self) -> Result<Option<Params>, NixlError> {
109        if !self.has_value()? {
110            return Ok(None);
111        }
112
113        let mut params = ptr::null_mut();
114        let status = unsafe {
115            nixl_capi_query_resp_list_get_params(self.list.inner.as_ptr(), self.index, &mut params)
116        };
117
118        match status {
119            NIXL_CAPI_SUCCESS => {
120                // SAFETY: If status is NIXL_CAPI_SUCCESS, params is non-null
121                let inner = unsafe { NonNull::new_unchecked(params) };
122                Ok(Some(Params::new(inner)))
123            }
124            NIXL_CAPI_ERROR_INVALID_PARAM => Err(NixlError::InvalidParam),
125            _ => Err(NixlError::BackendError),
126        }
127    }
128}
129
130/// An iterator over query responses
131pub struct QueryResponseIterator<'a> {
132    list: &'a QueryResponseList,
133    index: usize,
134    len: usize,
135}
136
137impl<'a> Iterator for QueryResponseIterator<'a> {
138    type Item = QueryResponse<'a>;
139
140    fn next(&mut self) -> Option<Self::Item> {
141        if self.index >= self.len {
142            None
143        } else {
144            let response = QueryResponse {
145                list: self.list,
146                index: self.index,
147            };
148            self.index += 1;
149            Some(response)
150        }
151    }
152
153    fn size_hint(&self) -> (usize, Option<usize>) {
154        let remaining = self.len - self.index;
155        (remaining, Some(remaining))
156    }
157}
158
159impl<'a> ExactSizeIterator for QueryResponseIterator<'a> {
160    fn len(&self) -> usize {
161        self.len - self.index
162    }
163}
164
165impl Drop for QueryResponseList {
166    fn drop(&mut self) {
167        // SAFETY: self.inner is guaranteed to be valid by NonNull
168        unsafe {
169            nixl_capi_destroy_query_resp_list(self.inner.as_ptr());
170        }
171    }
172}