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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

use std::collections::BTreeMap;

use crate::{
    error::Error,
    runtime::{dir::Dir, file::File},
    storage::types::Node,
};

const RESERVED_FD_COUNT: Fd = 3;

pub type Fd = u32;

pub enum FdEntry {
    File(File),
    Dir(Dir),
}

//
pub struct FdTable {
    // currently open file descriptors.
    table: BTreeMap<Fd, FdEntry>,
    // backward links to see how many file descriptors are currently pointing to any particular node.
    node_refcount: BTreeMap<Node, usize>,
    // the next generated descriptor's ID (if there is nothing to reuse).
    next_fd: Fd,
    // freed file descriptors ready to reuse.
    free_fds: Vec<Fd>,
}

impl FdTable {
    // create a new file descriptor table.
    pub fn new() -> Self {
        Self {
            table: BTreeMap::default(),
            node_refcount: BTreeMap::default(),
            next_fd: RESERVED_FD_COUNT,
            free_fds: vec![],
        }
    }

    // Get the map of node references.
    pub fn node_refcount(&self) -> &BTreeMap<Node, usize> {
        &self.node_refcount
    }

    // Update a file descriptor entry.
    pub fn update(&mut self, fd: Fd, entry: FdEntry) {
        self.insert(fd, entry);
    }

    // Update a file descriptor entry, it returns the old entry if existed.
    pub fn insert(&mut self, fd: Fd, entry: FdEntry) -> Option<FdEntry> {
        self.inc_node_refcount(&entry);

        let prev_entry = self.table.insert(fd, entry);
        if let Some(prev_entry) = prev_entry.as_ref() {
            self.dec_node_refcount(prev_entry);
        }

        prev_entry
    }

    // Get an FdEntry for a given file descriptor.
    pub fn get(&self, fd: Fd) -> Option<&FdEntry> {
        self.table.get(&fd)
    }

    // Open a new file descriptor.
    pub fn open(&mut self, entry: FdEntry) -> Fd {
        let fd = match self.free_fds.pop() {
            Some(fd) => fd,
            None => {
                let fd = self.next_fd;
                self.next_fd += 1;
                fd
            }
        };

        let prev = self.insert(fd, entry);
        assert!(prev.is_none());
        fd
    }

    // Reassign a file descriptor to a new number, the source descriptor is closed in the process.
    // If the destination descriptor is busy, it is closed in the process.
    pub fn renumber(&mut self, src: Fd, dst: Fd) -> Result<(), Error> {
        let old_entry = self.close(src).ok_or(Error::NotFound)?;

        // quietly close the destination file descriptor
        if let Some(_old_dst_entry) = self.close(dst) {
            let removed = self.free_fds.pop().unwrap();
            assert_eq!(removed, dst);
        }

        self.insert(dst, old_entry);

        Ok(())
    }

    // Close file descriptor.
    pub fn close(&mut self, fd: Fd) -> Option<FdEntry> {
        let entry = self.table.remove(&fd);

        if let Some(entry) = entry {
            self.free_fds.push(fd);
            self.dec_node_refcount(&entry);

            Some(entry)
        } else {
            None
        }
    }

    fn inc_node_refcount(&mut self, entry: &FdEntry) {
        let node = match entry {
            FdEntry::File(file) => file.node,
            FdEntry::Dir(dir) => dir.node,
        };
        let refcount = self.node_refcount.entry(node).or_default();
        *refcount += 1;
    }

    fn dec_node_refcount(&mut self, entry: &FdEntry) {
        let node = match entry {
            FdEntry::File(file) => file.node,
            FdEntry::Dir(dir) => dir.node,
        };

        let refcount = self.node_refcount.remove(&node);
        if let Some(mut refcount) = refcount {
            refcount -= 1;
            if refcount > 0 {
                self.node_refcount.insert(node, refcount);
            }
        }
    }
}