postgres_extension/
rust_utils.rs

1
2use std::cmp;
3use std::fmt;
4use std::io;
5use std::io::{Error,ErrorKind,Result};
6use std::mem;
7
8use crate::utils::memutils;
9use std::alloc::{GlobalAlloc, Layout};
10
11
12// implement Write trait for &[i8] (a.k.a. &[c_char])
13pub trait Write {
14    fn write(&mut self, buf: &[u8]) -> Result<usize>;
15    fn flush(&mut self) -> Result<()>;
16    fn write_all(&mut self, mut buf: &[u8]) -> Result<()> {
17        while !buf.is_empty() {
18            match self.write(buf) {
19                Ok(0) => return Err(Error::new(ErrorKind::WriteZero,
20                                               "failed to write whole buffer")),
21                Ok(n) => buf = &buf[n..],
22                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
23                Err(e) => return Err(e),
24            }
25        }
26        Ok(())
27    }
28    fn write_fmt(&mut self, fmt: fmt::Arguments) -> Result<()> {
29        // Create a shim which translates a Write to a fmt::Write and saves
30        // off I/O errors. instead of discarding them
31        struct Adaptor<'a, T: ?Sized + 'a> {
32            inner: &'a mut T,
33            error: Result<()>,
34        }
35
36        impl<'a, T: Write + ?Sized> fmt::Write for Adaptor<'a, T> {
37            fn write_str(&mut self, s: &str) -> fmt::Result {
38                match self.inner.write_all(s.as_bytes()) {
39                    Ok(()) => Ok(()),
40                    Err(e) => {
41                        self.error = Err(e);
42                        Err(fmt::Error)
43                    }
44                }
45            }
46        }
47
48        let mut output = Adaptor { inner: self, error: Ok(()) };
49        match fmt::write(&mut output, fmt) {
50            Ok(()) => Ok(()),
51            Err(..) => {
52                // check if the error came from the underlying `Write` or not
53                if output.error.is_err() {
54                    output.error
55                } else {
56                    Err(Error::new(ErrorKind::Other, "formatter error"))
57                }
58            }
59        }
60    }
61    fn by_ref(&mut self) -> &mut Self where Self: Sized { self }
62}
63
64impl<'a> Write for &'a mut [i8] {
65    #[inline]
66    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
67        let amt = cmp::min(data.len(), self.len());
68        let (a, b) = mem::replace(self, &mut []).split_at_mut(amt);
69        let a_u8 = unsafe { &mut *(a as *mut [i8] as *mut [u8]) };
70        a_u8.copy_from_slice(&data[..amt]);
71        *self = b;
72        Ok(amt)
73    }
74
75    #[inline]
76    fn write_all(&mut self, data: &[u8]) -> io::Result<()> {
77        if self.write(data)? == data.len() {
78            Ok(())
79        } else {
80            Err(Error::new(ErrorKind::WriteZero, "failed to write whole buffer"))
81        }
82    }
83
84    #[inline]
85    fn flush(&mut self) -> io::Result<()> { Ok(()) }
86}
87
88// Set up postgres allocator
89pub struct PostgresAllocator;
90unsafe impl GlobalAlloc for PostgresAllocator {
91    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
92        return memutils::c::MemoryContextAlloc(
93            memutils::CurrentMemoryContext, layout.size());
94    }
95    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
96        memutils::c::pfree(ptr);
97    }
98}
99
postgres_extension/rust_utils.rs

postgres_extension/
rust_utils.rs
