affine_importer 0.1.2

AFFiNE import planning and batching core.
Documentation
use std::collections::BTreeMap;

use super::{fragment_key, hash_bytes, is_protected, rewrite_markdown_lines};

#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct Wikilink<'a> {
  pub(super) raw: &'a str,
  pub(super) path: &'a str,
  pub(super) fragment: Option<&'a str>,
  pub(super) alias: Option<&'a str>,
  pub(super) embed: bool,
}

pub(super) fn rewrite_wikilinks(markdown: &str, mut replace: impl FnMut(Wikilink<'_>) -> String) -> String {
  rewrite_markdown_lines(markdown, |line, protected| {
    rewrite_wikilinks_in_line(line, protected, &mut replace)
  })
}

fn rewrite_wikilinks_in_line(
  line: &str,
  protected: &[std::ops::Range<usize>],
  replace: &mut impl FnMut(Wikilink<'_>) -> String,
) -> String {
  let mut output = String::with_capacity(line.len());
  let mut cursor = 0;
  while cursor < line.len() {
    let (embed, opener_len) = if line[cursor..].starts_with("![[") {
      (true, 3)
    } else if line[cursor..].starts_with("[[") {
      (false, 2)
    } else {
      let character = line[cursor..]
        .chars()
        .next()
        .expect("cursor is on a character boundary");
      output.push(character);
      cursor += character.len_utf8();
      continue;
    };
    if is_escaped(line, cursor) || is_protected(protected, cursor) {
      output.push_str(&line[cursor..cursor + opener_len]);
      cursor += opener_len;
      continue;
    }

    let content_start = cursor + opener_len;
    let Some(end_offset) = line[content_start..].find("]]") else {
      output.push_str(&line[cursor..]);
      break;
    };
    let end = content_start + end_offset;
    let raw_end = end + 2;
    let inner = &line[content_start..end];
    let (target, alias) = inner
      .split_once('|')
      .map_or((inner, None), |(target, alias)| (target, Some(alias)));
    let (path, fragment) = target
      .split_once('#')
      .map_or((target, None), |(path, fragment)| (path, Some(fragment)));
    output.push_str(&replace(Wikilink {
      raw: &line[cursor..raw_end],
      path: path.trim(),
      fragment: fragment.map(str::trim).filter(|fragment| !fragment.is_empty()),
      alias: alias.map(str::trim).filter(|alias| !alias.is_empty()),
      embed,
    }));
    cursor = raw_end;
  }
  output
}

pub(super) fn prepare_anchors(markdown: &str, doc_id: &str, hint_token: &str) -> (String, BTreeMap<String, String>) {
  let mut anchors = BTreeMap::new();
  let mut heading_path = Vec::<String>::new();
  let mut block_id_counts = BTreeMap::<String, usize>::new();
  let output = rewrite_markdown_lines(markdown, |line, protected| {
    let content_start = line.len() - line.trim_start().len();
    if is_protected(protected, content_start) {
      return line.to_string();
    }
    let (heading_line, explicit_anchor) = split_block_anchor(line)
      .map(|(content, anchor)| (content, Some(anchor)))
      .unwrap_or((line, None));
    if let Some((level, heading)) = parse_heading(heading_line) {
      heading_path.truncate(level.saturating_sub(1));
      heading_path.push(heading.to_string());
      let path = heading_path.join("#");
      let block_id = unique_block_id(
        explicit_anchor.map_or_else(
          || heading_block_id(doc_id, &path),
          |anchor| format!("block:{doc_id}:import:obsidian:{anchor}"),
        ),
        &mut block_id_counts,
      );
      anchors.entry(fragment_key(heading)).or_insert_with(|| block_id.clone());
      for start in 0..heading_path.len() {
        anchors
          .entry(fragment_key(&heading_path[start..].join("#")))
          .or_insert_with(|| block_id.clone());
      }
      if let Some(anchor) = explicit_anchor {
        anchors
          .entry(fragment_key(&format!("^{anchor}")))
          .or_insert_with(|| block_id.clone());
      }
      return format!("{heading_line} <!-- affine:block-id:{hint_token}={block_id} -->");
    }

    if let Some((content, anchor)) = split_block_anchor(line)
      && !is_protected(protected, line.rfind('^').unwrap_or(line.len()))
    {
      let block_id = unique_block_id(format!("block:{doc_id}:import:obsidian:{anchor}"), &mut block_id_counts);
      anchors
        .entry(fragment_key(&format!("^{anchor}")))
        .or_insert_with(|| block_id.clone());
      if content.trim().is_empty() {
        return format!("<!-- affine:block-id:{hint_token}={block_id} -->");
      } else {
        return format!("{content} <!-- affine:block-id:{hint_token}={block_id} -->");
      }
    }

    line.to_string()
  });

  (output, anchors)
}

pub(super) fn markdown_label(value: &str) -> String {
  let mut escaped = String::with_capacity(value.len());
  for character in value.chars() {
    if matches!(character, '[' | ']' | '\\') {
      escaped.push('\\');
    }
    escaped.push(character);
  }
  escaped
}

fn is_escaped(value: &str, index: usize) -> bool {
  value[..index].bytes().rev().take_while(|byte| *byte == b'\\').count() % 2 == 1
}

fn parse_heading(line: &str) -> Option<(usize, &str)> {
  let trimmed = line.trim_start();
  let level = trimmed.bytes().take_while(|byte| *byte == b'#').count();
  if !(1..=6).contains(&level) || trimmed.as_bytes().get(level) != Some(&b' ') {
    return None;
  }
  let heading = trimmed[level + 1..].trim().trim_end_matches('#').trim();
  (!heading.is_empty()).then_some((level, heading))
}

fn split_block_anchor(line: &str) -> Option<(&str, &str)> {
  let trimmed = line.trim_end();
  let start = trimmed
    .rfind(|character: char| character.is_whitespace())
    .map_or(0, |index| index + 1);
  let candidate = &trimmed[start..];
  let anchor = candidate.strip_prefix('^')?;
  if anchor.is_empty()
    || !anchor
      .chars()
      .all(|character| character.is_ascii_alphanumeric() || character == '-')
  {
    return None;
  }
  Some((trimmed[..start].trim_end(), anchor))
}

fn heading_block_id(doc_id: &str, heading: &str) -> String {
  let hash = hash_bytes(heading.to_lowercase().as_bytes());
  format!("block:{doc_id}:import:obsidian:heading:{}", &hash[..16])
}

fn unique_block_id(base: String, counts: &mut BTreeMap<String, usize>) -> String {
  let count = counts.entry(base.clone()).or_default();
  let id = if *count == 0 {
    base
  } else {
    format!("{base}:{}", *count + 1)
  };
  *count += 1;
  id
}

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn prepares_anchors_without_rewriting_code() {
    let (markdown, anchors) = prepare_anchors(
      "# Title\n\n## Parent\n\n### Child\n\n## Parent\n\nText ^anchor\n\n- Item ^list-item\n\n```md\n## \
       Code\n[[Target]]\ntext ^code\n```\n\n    ## Indented\n    text ^indented",
      "doc",
      "token",
    );

    assert!(markdown.contains("affine:block-id:token=block:doc:import:obsidian:anchor"));
    assert!(markdown.contains("affine:block-id:token=block:doc:import:obsidian:list-item"));
    assert!(anchors.contains_key("title"));
    assert!(anchors.contains_key("parent#child"));
    assert!(!anchors.contains_key("code"));
    assert!(!anchors.contains_key("^code"));
    assert!(!anchors.contains_key("indented"));
    assert!(!anchors.contains_key("^indented"));
    assert_eq!(anchors["^anchor"], "block:doc:import:obsidian:anchor");
    let heading_ids = markdown
      .lines()
      .filter(|line| line.starts_with("## Parent"))
      .collect::<Vec<_>>();
    assert_ne!(heading_ids[0], heading_ids[1]);

    let rewritten = rewrite_wikilinks("`[[Inline]]`\n```\n[[Fence]]\n```\n[[Page]]", |_| "LINK".to_string());
    assert_eq!(rewritten, "`[[Inline]]`\n```\n[[Fence]]\n```\nLINK");
    let rewritten = rewrite_wikilinks("`start\n[[Multiline]]\nend`\n\n    [[Indented]]\n\n[[Page]]", |_| {
      "LINK".to_string()
    });
    assert_eq!(rewritten, "`start\n[[Multiline]]\nend`\n\n    [[Indented]]\n\nLINK");
  }
}