atmos/kernel/fs/

btree.rs

// btree.rs - Disk-based Copy-on-Write (COW) B-Tree for indexing filenames and labels
#![allow(dead_code)]
#![allow(static_mut_refs)]

extern crate alloc;
use alloc::vec::Vec;

use super::block::{BlockDevice, MirroredBlockDevice, DATA_SIZE_PER_SECTOR};

const MAX_KEYS: usize = 24;
const KEY_LEN: usize = 16; // 128-bit hash

/// ディスク上のB-Treeノード (データサイズ 508 バイト = セクタデータ領域ちょうど)。
/// 挿入は「挿入→満杯なら分割」方式のため num_keys は分割直前に一時的に MAX_KEYS+1 まで増える。
/// その一時オーバーフローを格納できるよう keys は MAX_KEYS+1、values は内部ノードの
/// スライド (values[i+2]) と最右子ポインタを賄うため MAX_KEYS+2 を確保する。
/// 永続化されるノードの num_keys は分割により常に MAX_KEYS 以下。
#[repr(C, packed)]
#[derive(Clone, Copy)]
pub struct BTreeNode {
    pub node_type: u8, // 0: Leaf, 1: Internal
    pub num_keys: u8,
    pub padding: [u8; 2],
    pub keys: [[u8; KEY_LEN]; MAX_KEYS + 1], // UTF-8 キー配列 (+1 = 分割前の一時オーバーフロー)
    pub values: [u32; MAX_KEYS + 2], // 葉:Metadataセクタ / 内部:子ノードセクタ (+2 = 一時オーバーフロー)
}

// BTreeNode は 1 セクタのデータ領域 (508B) にちょうど収まる必要がある。
// keys/values の容量を変えてここを超えると to_bytes/from_bytes が領域外アクセスするため、
// コンパイル時に検知する。
const _: () = assert!(core::mem::size_of::<BTreeNode>() <= DATA_SIZE_PER_SECTOR);

impl BTreeNode {
    pub fn new(is_leaf: bool) -> Self {
        Self {
            node_type: if is_leaf { 0 } else { 1 },
            num_keys: 0,
            padding: [0; 2],
            keys: [[0; KEY_LEN]; MAX_KEYS + 1],
            values: [0; MAX_KEYS + 2],
        }
    }

    /// バイトスライスから復元
    pub fn from_bytes(bytes: &[u8]) -> Self {
        let mut node = Self::new(true);
        let ptr = &mut node as *mut Self as *mut u8;
        unsafe {
            core::ptr::copy_nonoverlapping(bytes.as_ptr(), ptr, core::mem::size_of::<Self>());
        }
        node
    }

    /// バイトスライスへシリアライズ
    pub fn to_bytes(&self) -> [u8; DATA_SIZE_PER_SECTOR] {
        let mut buf = [0u8; DATA_SIZE_PER_SECTOR];
        let ptr = self as *const Self as *const u8;
        unsafe {
            core::ptr::copy_nonoverlapping(ptr, buf.as_mut_ptr(), core::mem::size_of::<Self>());
        }
        buf
    }

    /// キー文字列の比較ヘルパー (ハッシュ値の単純比較)
    fn compare_keys(a: &[u8; KEY_LEN], b: &[u8]) -> core::cmp::Ordering {
        let b_len = b.len().min(KEY_LEN);
        let mut b_padded = [0u8; KEY_LEN];
        b_padded[..b_len].copy_from_slice(&b[..b_len]);
        a.cmp(&b_padded)
    }
}

/// ディスク上のCOW B-Treeインデックスマネージャ
pub struct CowBTree {
    root_sector: u32,
}

impl CowBTree {
    pub fn new(root_sector: u32) -> Self {
        Self { root_sector }
    }

    pub fn get_root(&self) -> u32 {
        self.root_sector
    }

    /// 指定したキーを B-Tree から検索します (O(log N) 探索)
    /// 成功すると、指し示されているファイルメタデータセクタ番号を返します。
    pub fn search<B: BlockDevice>(
        &self,
        device: &mut MirroredBlockDevice<B>,
        key: &[u8; KEY_LEN],
    ) -> Option<u32> {
        let mut curr_sector = self.root_sector;
        let mut buf = [0u8; DATA_SIZE_PER_SECTOR];
        let mut loop_count = 0;

        while curr_sector != 0 {
            loop_count += 1;
            if loop_count > 100 {
                crate::println!(
                    "      [BTree DBG] search loop exceeded 100 iterations! curr_sector={}",
                    curr_sector
                );
                return None;
            }
            if device.read_sector(curr_sector as usize, &mut buf).is_err() {
                return None;
            }
            let node = BTreeNode::from_bytes(&buf);

            // 現在のノード内を探索
            let mut idx = 0;
            while idx < node.num_keys as usize {
                if idx >= MAX_KEYS {
                    crate::println!(
                        "      [BTree DBG] search node num_keys is corrupted! num_keys={}",
                        node.num_keys
                    );
                    return None;
                }
                let cmp = BTreeNode::compare_keys(&node.keys[idx], key);
                if cmp == core::cmp::Ordering::Equal {
                    if node.node_type == 0 {
                        // 葉ノード: メタデータセクタ番号を返す
                        return Some(node.values[idx]);
                    } else {
                        // 内部ノード: 次の子ノードへ遷移
                        idx += 1;
                        break;
                    }
                } else if cmp == core::cmp::Ordering::Greater {
                    break;
                }
                idx += 1;
            }

            if node.node_type == 0 {
                break; // 葉ノードで見つからなかった
            } else {
                curr_sector = node.values[idx]; // 子ノードセクタへ下る
            }
        }

        None
    }

    /// インデックスへ新しいキーとセクタのペアを挿入します (COWによる再帰的ツリー更新)
    /// 更新された新しいルートノードのセクタ番号を返します。
    pub fn insert<B: BlockDevice>(
        &mut self,
        device: &mut MirroredBlockDevice<B>,
        key: &[u8; KEY_LEN],
        val_sector: u32,
        alloc_sector_fn: &mut dyn FnMut() -> u32,
    ) -> Result<u32, &'static str> {
        let root_sec = self.root_sector;

        // 1. ルートがまだ存在しない場合、新しい葉ノードを作成
        if root_sec == 0 {
            let new_sec = alloc_sector_fn();
            if new_sec == 0 {
                return Err("btree: sector allocation failed");
            }

            let mut new_node = BTreeNode::new(true);
            new_node.num_keys = 1;
            new_node.keys[0] = *key;
            new_node.values[0] = val_sector;

            device.write_sector(new_sec as usize, &new_node.to_bytes())?;
            self.root_sector = new_sec;
            return Ok(new_sec);
        }

        // 2. COW 再帰挿入の実行
        let result = self.insert_recursive(device, root_sec, key, val_sector, alloc_sector_fn)?;

        match result {
            InsertResult::Ok(new_sec) => {
                self.root_sector = new_sec;
                Ok(new_sec)
            }
            InsertResult::Split(new_sec, split_key, split_sibling_sec) => {
                // ルートノードが分割された場合、新しいルート（内部ノード）を作成
                let new_root_sec = alloc_sector_fn();
                if new_root_sec == 0 {
                    return Err("btree: sector allocation failed");
                }

                let mut new_root = BTreeNode::new(false);
                new_root.num_keys = 1;
                new_root.keys[0] = split_key;
                new_root.values[0] = new_sec;
                new_root.values[1] = split_sibling_sec;

                device.write_sector(new_root_sec as usize, &new_root.to_bytes())?;
                self.root_sector = new_root_sec;
                Ok(new_root_sec)
            }
        }
    }

    /// すべての葉ノードエントリ (key, value_sector) を収集
    pub fn collect_leaf_entries<B: BlockDevice>(
        &self,
        device: &mut MirroredBlockDevice<B>,
        out: &mut Vec<([u8; KEY_LEN], u32)>,
    ) {
        if self.root_sector == 0 {
            return;
        }
        let mut visited = Vec::new();
        self.collect_from_node(device, self.root_sector, out, &mut visited);
    }

    fn collect_from_node<B: BlockDevice>(
        &self,
        device: &mut MirroredBlockDevice<B>,
        sector: u32,
        out: &mut Vec<([u8; KEY_LEN], u32)>,
        visited: &mut Vec<u32>,
    ) {
        if visited.contains(&sector) {
            crate::println!(
                "      [BTree DBG] Circular reference detected in collect_from_node: sector={}",
                sector
            );
            return;
        }
        visited.push(sector);
        if visited.len() > 100 {
            crate::println!("      [BTree DBG] Recursion limit exceeded in collect_from_node!");
            return;
        }

        let mut buf = [0u8; DATA_SIZE_PER_SECTOR];
        if device.read_sector(sector as usize, &mut buf).is_err() {
            return;
        }

        let node = BTreeNode::from_bytes(&buf);
        if node.node_type == 0 {
            let limit = (node.num_keys as usize).min(MAX_KEYS);
            for i in 0..limit {
                out.push((node.keys[i], node.values[i]));
            }
            return;
        }

        let limit = (node.num_keys as usize).min(MAX_KEYS);
        for i in 0..=limit {
            let child = node.values[i];
            if child != 0 {
                self.collect_from_node(device, child, out, visited);
            }
        }
    }

    /// COW を保ちつつ、B-Treeに再帰的に挿入を行います。
    fn insert_recursive<B: BlockDevice>(
        &self,
        device: &mut MirroredBlockDevice<B>,
        curr_sector: u32,
        key: &[u8; KEY_LEN],
        val_sector: u32,
        alloc_sector_fn: &mut dyn FnMut() -> u32,
    ) -> Result<InsertResult, &'static str> {
        let mut buf = [0u8; DATA_SIZE_PER_SECTOR];
        // 【2026-08-05】ここは `?` で `read_sector` のエラーをそのまま返していた。
        //
        // 実測で `[CACHE] 保存に失敗 ... err=read_sector: sector out of range`
        // が出続けており、5.3MB のフォントも 8KB の JS も保存できていない
        // （＝毎回 14MB を取り直している）。だが原文のエラーだけでは
        // **どのセクタが範囲外なのか、どこから来た番号なのか**が分からず、
        // 原因に辿り着けなかった。
        //
        // 木の節点番号が実在範囲を外れているのは索引の破損なので、
        // 値と範囲を添えて報告する（推測で直さないための情報を残す）。
        if curr_sector as usize >= device.total_usable_sectors() {
            crate::warn!(
                "[SylFS][BTREE] 節点番号が実在範囲外 curr_sector={} usable={} (索引が壊れている)",
                curr_sector,
                device.total_usable_sectors()
            );
            return Err("btree: node sector out of range");
        }
        if let Err(e) = device.read_sector(curr_sector as usize, &mut buf) {
            crate::warn!(
                "[SylFS][BTREE] 節点の読み出しに失敗 curr_sector={} usable={} err={}",
                curr_sector,
                device.total_usable_sectors(),
                e
            );
            return Err(e);
        }

        // COW: 読み込んだノードを新しいセクタにコピーして書き換える準備をする
        let mut node = BTreeNode::from_bytes(&buf);
        let new_sector = alloc_sector_fn();
        if new_sector == 0 {
            return Err("btree: sector allocation failed");
        }

        // 現在のノード内での挿入位置を二分探索/線形探索
        let mut idx = 0;
        while idx < node.num_keys as usize {
            let cmp = BTreeNode::compare_keys(&node.keys[idx], key);
            if cmp == core::cmp::Ordering::Equal {
                if node.node_type == 0 {
                    // 葉ノード: すでに同じキーが存在 → 値(メタデータセクタ)を上書き更新
                    node.values[idx] = val_sector;
                    device.write_sector(new_sector as usize, &node.to_bytes())?;
                    return Ok(InsertResult::Ok(new_sector));
                } else {
                    // 内部ノード: 分離キーと一致した場合、実体は右部分木の葉にある（B+Tree）。
                    // search と同様に右の子へ降下する（values[idx] を上書きすると部分木を破壊する）。
                    idx += 1;
                    break;
                }
            } else if cmp == core::cmp::Ordering::Greater {
                break;
            }
            idx += 1;
        }

        if node.node_type == 0 {
            // --- 葉ノード (Leaf) への直接挿入 ---
            // キーと値をスライドさせて挿入
            for i in (idx..node.num_keys as usize).rev() {
                node.keys[i + 1] = node.keys[i];
                node.values[i + 1] = node.values[i];
            }
            node.keys[idx] = *key;
            node.values[idx] = val_sector;
            node.num_keys += 1;

            if node.num_keys as usize > MAX_KEYS {
                // 分割 (Split) が発生
                let split_res = self.split_node(device, &mut node, new_sector, alloc_sector_fn)?;
                Ok(split_res)
            } else {
                // 分割なし
                device.write_sector(new_sector as usize, &node.to_bytes())?;
                Ok(InsertResult::Ok(new_sector))
            }
        } else {
            // --- 内部ノード (Internal) への再帰的挿入 ---
            let child_sec = node.values[idx];
            let sub_result =
                self.insert_recursive(device, child_sec, key, val_sector, alloc_sector_fn)?;

            match sub_result {
                InsertResult::Ok(new_child_sec) => {
                    // 子ノードがCOWで新しいセクタに切り替わったため、ポインタを更新
                    node.values[idx] = new_child_sec;
                    device.write_sector(new_sector as usize, &node.to_bytes())?;
                    Ok(InsertResult::Ok(new_sector))
                }
                InsertResult::Split(new_child_sec, split_key, split_sibling_sec) => {
                    // 子ノードが分割されたため、分割キーと新しい兄弟へのポインタを親ノード(自分)に挿入
                    node.values[idx] = new_child_sec;

                    for i in (idx..node.num_keys as usize).rev() {
                        node.keys[i + 1] = node.keys[i];
                        node.values[i + 2] = node.values[i + 1];
                    }
                    node.keys[idx] = split_key;
                    node.values[idx + 1] = split_sibling_sec;
                    node.num_keys += 1;

                    if node.num_keys as usize > MAX_KEYS {
                        // 親ノードもさらに分割
                        let split_res =
                            self.split_node(device, &mut node, new_sector, alloc_sector_fn)?;
                        Ok(split_res)
                    } else {
                        device.write_sector(new_sector as usize, &node.to_bytes())?;
                        Ok(InsertResult::Ok(new_sector))
                    }
                }
            }
        }
    }

    /// ノードの半分を別の新しいセクタに分割し、親への引き上げ情報を返します (COW対応)
    fn split_node<B: BlockDevice>(
        &self,
        device: &mut MirroredBlockDevice<B>,
        node: &mut BTreeNode,
        new_sector: u32,
        alloc_sector_fn: &mut dyn FnMut() -> u32,
    ) -> Result<InsertResult, &'static str> {
        let sibling_sector = alloc_sector_fn();
        if sibling_sector == 0 {
            return Err("btree: sector allocation failed");
        }

        let mut sibling = BTreeNode::new(node.node_type == 0);
        // 実際の num_keys を基準に分割する (一時オーバーフロー時は MAX_KEYS+1)。
        let n = node.num_keys as usize;
        let mid = n / 2;

        // 親に引き上げる分離キー
        let split_key = node.keys[mid];

        if node.node_type == 0 {
            // --- 葉ノード (B+Tree): 中央キーは右兄弟に残し、分離キーとしてコピー引き上げ。
            //     キーを取りこぼさないよう右兄弟は keys[mid..n] を受け取る。
            let sibling_keys_num = n - mid;
            for i in 0..sibling_keys_num {
                sibling.keys[i] = node.keys[mid + i];
                sibling.values[i] = node.values[mid + i];
            }
            sibling.num_keys = sibling_keys_num as u8;
        } else {
            // --- 内部ノード (B-Tree): 中央キーは親へ移動し子からは除去。
            //     右兄弟は keys[mid+1..n] と対応する子ポインタ values[mid+1..n+1] を受け取る。
            let sibling_keys_num = n - mid - 1;
            for i in 0..sibling_keys_num {
                sibling.keys[i] = node.keys[mid + 1 + i];
                sibling.values[i] = node.values[mid + 1 + i];
            }
            sibling.values[sibling_keys_num] = node.values[n]; // 最右子ポインタ
            sibling.num_keys = sibling_keys_num as u8;
        }

        // 前半のノードサイズを縮小 (左ノードは keys[0..mid] を保持)
        node.num_keys = mid as u8;

        // 両ノードをそれぞれ別の空きセクタに書き込み (COW)
        device.write_sector(new_sector as usize, &node.to_bytes())?;
        device.write_sector(sibling_sector as usize, &sibling.to_bytes())?;

        Ok(InsertResult::Split(new_sector, split_key, sibling_sector))
    }
}

/// 挿入再帰処理の結果
enum InsertResult {
    /// 通常終了。新しいセクタ番号を返す。
    Ok(u32),
    /// 分割発生。新しいセクタ番号、引き上げられるキー、および新しい右側の兄弟セクタを返す。
    Split(u32, [u8; KEY_LEN], u32),
}