use std::collections::{HashMap, HashSet};
use std::fmt;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LinearKind {
Unrestricted,
Linear,
Affine,
Relevant,
}
impl LinearKind {
pub fn allows_copy(&self) -> bool {
matches!(self, LinearKind::Unrestricted | LinearKind::Relevant)
}
pub fn allows_drop(&self) -> bool {
matches!(self, LinearKind::Unrestricted | LinearKind::Affine)
}
pub fn requires_use(&self) -> bool {
matches!(self, LinearKind::Linear | LinearKind::Relevant)
}
pub fn limits_use(&self) -> bool {
matches!(self, LinearKind::Linear | LinearKind::Affine)
}
pub fn join(&self, other: &LinearKind) -> LinearKind {
match (*self, *other) {
(LinearKind::Unrestricted, _) | (_, LinearKind::Unrestricted) => {
LinearKind::Unrestricted
}
(LinearKind::Linear, LinearKind::Linear) => LinearKind::Linear,
(LinearKind::Affine, LinearKind::Affine) => LinearKind::Affine,
(LinearKind::Relevant, LinearKind::Relevant) => LinearKind::Relevant,
(LinearKind::Linear, LinearKind::Affine) | (LinearKind::Affine, LinearKind::Linear) => {
LinearKind::Affine
}
(LinearKind::Linear, LinearKind::Relevant)
| (LinearKind::Relevant, LinearKind::Linear) => LinearKind::Relevant,
(LinearKind::Affine, LinearKind::Relevant)
| (LinearKind::Relevant, LinearKind::Affine) => LinearKind::Unrestricted,
}
}
pub fn meet(&self, other: &LinearKind) -> LinearKind {
match (*self, *other) {
(LinearKind::Linear, _) | (_, LinearKind::Linear) => LinearKind::Linear,
(LinearKind::Affine, LinearKind::Affine) => LinearKind::Affine,
(LinearKind::Relevant, LinearKind::Relevant) => LinearKind::Relevant,
(LinearKind::Affine, LinearKind::Relevant)
| (LinearKind::Relevant, LinearKind::Affine) => LinearKind::Linear,
(LinearKind::Unrestricted, other) | (other, LinearKind::Unrestricted) => other,
}
}
}
impl fmt::Display for LinearKind {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
LinearKind::Unrestricted => write!(f, "unrestricted"),
LinearKind::Linear => write!(f, "linear"),
LinearKind::Affine => write!(f, "affine"),
LinearKind::Relevant => write!(f, "relevant"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Ownership {
Owned,
Moved,
Borrowed,
Dropped,
}
impl fmt::Display for Ownership {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Ownership::Owned => write!(f, "owned"),
Ownership::Moved => write!(f, "moved"),
Ownership::Borrowed => write!(f, "borrowed"),
Ownership::Dropped => write!(f, "dropped"),
}
}
}
#[derive(Debug, Clone)]
pub struct LinearType {
pub base_type: String,
pub kind: LinearKind,
pub name: Option<String>,
pub description: Option<String>,
pub tags: Vec<String>,
}
impl LinearType {
pub fn new(base_type: impl Into<String>) -> Self {
LinearType {
base_type: base_type.into(),
kind: LinearKind::Unrestricted,
name: None,
description: None,
tags: Vec::new(),
}
}
pub fn linear(base_type: impl Into<String>) -> Self {
LinearType {
base_type: base_type.into(),
kind: LinearKind::Linear,
name: None,
description: None,
tags: Vec::new(),
}
}
pub fn affine(base_type: impl Into<String>) -> Self {
LinearType {
base_type: base_type.into(),
kind: LinearKind::Affine,
name: None,
description: None,
tags: Vec::new(),
}
}
pub fn relevant(base_type: impl Into<String>) -> Self {
LinearType {
base_type: base_type.into(),
kind: LinearKind::Relevant,
name: None,
description: None,
tags: Vec::new(),
}
}
pub fn with_kind(mut self, kind: LinearKind) -> Self {
self.kind = kind;
self
}
pub fn with_name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
pub fn with_description(mut self, description: impl Into<String>) -> Self {
self.description = Some(description.into());
self
}
pub fn with_tag(mut self, tag: impl Into<String>) -> Self {
self.tags.push(tag.into());
self
}
pub fn type_name(&self) -> &str {
self.name.as_deref().unwrap_or(&self.base_type)
}
pub fn allows_copy(&self) -> bool {
self.kind.allows_copy()
}
pub fn allows_drop(&self) -> bool {
self.kind.allows_drop()
}
}
#[derive(Debug, Clone)]
pub struct Resource {
pub name: String,
pub ty: LinearType,
pub ownership: Ownership,
pub use_count: usize,
pub created_at: Option<String>,
pub last_used_at: Option<String>,
}
impl Resource {
pub fn new(name: impl Into<String>, ty: LinearType) -> Self {
Resource {
name: name.into(),
ty,
ownership: Ownership::Owned,
use_count: 0,
created_at: None,
last_used_at: None,
}
}
pub fn with_created_at(mut self, location: impl Into<String>) -> Self {
self.created_at = Some(location.into());
self
}
pub fn can_use(&self) -> bool {
matches!(self.ownership, Ownership::Owned)
|| (matches!(self.ownership, Ownership::Borrowed) && self.ty.kind.allows_copy())
}
pub fn can_move(&self) -> bool {
matches!(self.ownership, Ownership::Owned)
}
pub fn can_drop(&self) -> bool {
self.ty.allows_drop() || self.use_count > 0
}
pub fn use_resource(&mut self, location: impl Into<String>) -> Result<(), LinearError> {
if !self.can_use() {
return Err(LinearError::UseAfterMove {
resource: self.name.clone(),
state: self.ownership,
});
}
if self.ty.kind.limits_use() && self.use_count > 0 {
return Err(LinearError::MultipleUse {
resource: self.name.clone(),
count: self.use_count + 1,
});
}
self.use_count += 1;
self.last_used_at = Some(location.into());
Ok(())
}
pub fn move_to(&mut self, location: impl Into<String>) -> Result<(), LinearError> {
if !self.can_move() {
return Err(LinearError::UseAfterMove {
resource: self.name.clone(),
state: self.ownership,
});
}
self.ownership = Ownership::Moved;
self.last_used_at = Some(location.into());
Ok(())
}
pub fn drop_resource(&mut self) -> Result<(), LinearError> {
if self.ty.kind.requires_use() && self.use_count == 0 {
return Err(LinearError::UnusedResource {
resource: self.name.clone(),
kind: self.ty.kind,
});
}
self.ownership = Ownership::Dropped;
Ok(())
}
pub fn validate_end_of_scope(&self) -> Result<(), LinearError> {
match self.ownership {
Ownership::Owned => {
if self.ty.kind.requires_use() && self.use_count == 0 {
Err(LinearError::UnusedResource {
resource: self.name.clone(),
kind: self.ty.kind,
})
} else {
Ok(())
}
}
Ownership::Borrowed => Err(LinearError::BorrowedAtEndOfScope {
resource: self.name.clone(),
}),
_ => Ok(()),
}
}
}
#[derive(Debug, Clone)]
pub enum LinearError {
UseAfterMove { resource: String, state: Ownership },
MultipleUse { resource: String, count: usize },
UnusedResource { resource: String, kind: LinearKind },
BorrowedAtEndOfScope { resource: String },
UnknownResource { resource: String },
DuplicateResource { resource: String },
}
impl fmt::Display for LinearError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
LinearError::UseAfterMove { resource, state } => {
write!(f, "Cannot use resource '{}': it is {}", resource, state)
}
LinearError::MultipleUse { resource, count } => {
write!(
f,
"Resource '{}' used {} times but must be used exactly once",
resource, count
)
}
LinearError::UnusedResource { resource, kind } => {
write!(f, "{} resource '{}' was never used", kind, resource)
}
LinearError::BorrowedAtEndOfScope { resource } => {
write!(
f,
"Resource '{}' is still borrowed at end of scope",
resource
)
}
LinearError::UnknownResource { resource } => {
write!(f, "Unknown resource '{}'", resource)
}
LinearError::DuplicateResource { resource } => {
write!(f, "Resource '{}' already exists", resource)
}
}
}
}
impl std::error::Error for LinearError {}
#[derive(Debug, Clone, Default)]
pub struct LinearContext {
resources: HashMap<String, Resource>,
aliases: HashMap<String, String>,
scope_stack: Vec<HashSet<String>>,
}
impl LinearContext {
pub fn new() -> Self {
LinearContext {
resources: HashMap::new(),
aliases: HashMap::new(),
scope_stack: vec![HashSet::new()],
}
}
pub fn enter_scope(&mut self) {
self.scope_stack.push(HashSet::new());
}
pub fn exit_scope(&mut self) -> Result<(), Vec<LinearError>> {
let scope = match self.scope_stack.pop() {
Some(s) => s,
None => return Ok(()),
};
let mut errors = Vec::new();
for resource_name in scope {
if let Some(resource) = self.resources.get(&resource_name) {
if let Err(e) = resource.validate_end_of_scope() {
errors.push(e);
}
}
}
for resource_name in self.scope_stack.last().cloned().unwrap_or_default() {
self.resources.remove(&resource_name);
}
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
pub fn add_resource(&mut self, resource: Resource) -> Result<(), LinearError> {
if self.resources.contains_key(&resource.name) {
return Err(LinearError::DuplicateResource {
resource: resource.name,
});
}
let name = resource.name.clone();
self.resources.insert(name.clone(), resource);
if let Some(scope) = self.scope_stack.last_mut() {
scope.insert(name);
}
Ok(())
}
pub fn create_resource(
&mut self,
name: impl Into<String>,
ty: LinearType,
location: impl Into<String>,
) -> Result<(), LinearError> {
let resource = Resource::new(name, ty).with_created_at(location);
self.add_resource(resource)
}
pub fn get_resource(&self, name: &str) -> Option<&Resource> {
self.resolve_alias(name).and_then(|n| self.resources.get(n))
}
pub fn get_resource_mut(&mut self, name: &str) -> Option<&mut Resource> {
let resolved = self.resolve_alias(name).map(|s| s.to_string());
resolved.and_then(move |n| self.resources.get_mut(&n))
}
fn resolve_alias<'a>(&'a self, name: &'a str) -> Option<&'a str> {
if self.resources.contains_key(name) {
return Some(name);
}
self.aliases.get(name).map(|s| s.as_str())
}
pub fn use_resource(
&mut self,
name: &str,
location: impl Into<String>,
) -> Result<(), LinearError> {
let resource = self
.get_resource_mut(name)
.ok_or_else(|| LinearError::UnknownResource {
resource: name.to_string(),
})?;
resource.use_resource(location)
}
pub fn move_resource(
&mut self,
from: &str,
to: impl Into<String>,
location: impl Into<String>,
) -> Result<(), LinearError> {
let to = to.into();
let location = location.into();
let from_resolved = self
.resolve_alias(from)
.ok_or_else(|| LinearError::UnknownResource {
resource: from.to_string(),
})?
.to_string();
let resource =
self.resources
.get_mut(&from_resolved)
.ok_or_else(|| LinearError::UnknownResource {
resource: from.to_string(),
})?;
resource.move_to(&location)?;
self.aliases.insert(to.clone(), from_resolved.clone());
if let Some(scope) = self.scope_stack.last_mut() {
scope.insert(to);
}
Ok(())
}
pub fn drop_resource(&mut self, name: &str) -> Result<(), LinearError> {
let resource = self
.get_resource_mut(name)
.ok_or_else(|| LinearError::UnknownResource {
resource: name.to_string(),
})?;
resource.drop_resource()
}
pub fn validate_all(&self) -> Result<(), Vec<LinearError>> {
let mut errors = Vec::new();
for resource in self.resources.values() {
if let Err(e) = resource.validate_end_of_scope() {
errors.push(e);
}
}
if errors.is_empty() {
Ok(())
} else {
Err(errors)
}
}
pub fn resource_names(&self) -> Vec<&str> {
self.resources.keys().map(|s| s.as_str()).collect()
}
pub fn len(&self) -> usize {
self.resources.len()
}
pub fn is_empty(&self) -> bool {
self.resources.is_empty()
}
pub fn statistics(&self) -> LinearStatistics {
let mut total = 0;
let mut used = 0;
let mut unused = 0;
let mut moved = 0;
for resource in self.resources.values() {
total += 1;
match resource.ownership {
Ownership::Moved => moved += 1,
_ => {
if resource.use_count > 0 {
used += 1;
} else {
unused += 1;
}
}
}
}
LinearStatistics {
total,
used,
unused,
moved,
}
}
}
#[derive(Debug, Clone)]
pub struct LinearStatistics {
pub total: usize,
pub used: usize,
pub unused: usize,
pub moved: usize,
}
#[derive(Debug, Clone, Default)]
pub struct LinearTypeRegistry {
types: HashMap<String, LinearType>,
}
impl LinearTypeRegistry {
pub fn new() -> Self {
LinearTypeRegistry {
types: HashMap::new(),
}
}
pub fn with_builtins() -> Self {
let mut registry = LinearTypeRegistry::new();
registry.register(
LinearType::linear("Tensor")
.with_name("GpuTensor")
.with_tag("gpu")
.with_description("GPU tensor that must be explicitly freed"),
);
registry.register(
LinearType::affine("FileHandle")
.with_name("FileHandle")
.with_tag("io")
.with_description("File handle that can be closed or dropped"),
);
registry.register(
LinearType::linear("Connection")
.with_name("NetworkConnection")
.with_tag("network")
.with_description("Network connection that must be closed"),
);
registry.register(
LinearType::linear("Guard")
.with_name("MutexGuard")
.with_tag("sync")
.with_description("Mutex guard that must be released"),
);
registry
}
pub fn register(&mut self, ty: LinearType) {
let name = ty.type_name().to_string();
self.types.insert(name, ty);
}
pub fn get(&self, name: &str) -> Option<&LinearType> {
self.types.get(name)
}
pub fn contains(&self, name: &str) -> bool {
self.types.contains_key(name)
}
pub fn type_names(&self) -> Vec<&str> {
self.types.keys().map(|s| s.as_str()).collect()
}
pub fn len(&self) -> usize {
self.types.len()
}
pub fn is_empty(&self) -> bool {
self.types.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_linear_kind_properties() {
assert!(!LinearKind::Linear.allows_copy());
assert!(!LinearKind::Linear.allows_drop());
assert!(!LinearKind::Affine.allows_copy());
assert!(LinearKind::Affine.allows_drop());
assert!(LinearKind::Relevant.allows_copy());
assert!(!LinearKind::Relevant.allows_drop());
assert!(LinearKind::Unrestricted.allows_copy());
assert!(LinearKind::Unrestricted.allows_drop());
}
#[test]
fn test_linear_kind_join() {
assert_eq!(
LinearKind::Linear.join(&LinearKind::Linear),
LinearKind::Linear
);
assert_eq!(
LinearKind::Affine.join(&LinearKind::Relevant),
LinearKind::Unrestricted
);
}
#[test]
fn test_linear_kind_meet() {
assert_eq!(
LinearKind::Affine.meet(&LinearKind::Relevant),
LinearKind::Linear
);
assert_eq!(
LinearKind::Linear.meet(&LinearKind::Affine),
LinearKind::Linear
);
}
#[test]
fn test_create_resource() {
let mut ctx = LinearContext::new();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
assert!(ctx.get_resource("x").is_some());
}
#[test]
fn test_use_linear_resource_once() {
let mut ctx = LinearContext::new();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
ctx.use_resource("x", "line 5").expect("unwrap");
let resource = ctx.get_resource("x").expect("unwrap");
assert_eq!(resource.use_count, 1);
}
#[test]
fn test_use_linear_resource_twice_fails() {
let mut ctx = LinearContext::new();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
ctx.use_resource("x", "line 5").expect("unwrap");
let result = ctx.use_resource("x", "line 10");
assert!(matches!(result, Err(LinearError::MultipleUse { .. })));
}
#[test]
fn test_affine_can_be_dropped() {
let mut ctx = LinearContext::new();
let ty = LinearType::affine("Handle");
ctx.create_resource("h", ty, "line 1").expect("unwrap");
ctx.drop_resource("h").expect("unwrap");
}
#[test]
fn test_linear_unused_fails() {
let mut ctx = LinearContext::new();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
let result = ctx.validate_all();
assert!(result.is_err());
}
#[test]
fn test_relevant_can_be_used_multiple_times() {
let mut ctx = LinearContext::new();
let ty = LinearType::relevant("Value");
ctx.create_resource("v", ty, "line 1").expect("unwrap");
ctx.use_resource("v", "line 5").expect("unwrap");
ctx.use_resource("v", "line 10").expect("unwrap");
ctx.use_resource("v", "line 15").expect("unwrap");
let resource = ctx.get_resource("v").expect("unwrap");
assert_eq!(resource.use_count, 3);
}
#[test]
fn test_move_resource() {
let mut ctx = LinearContext::new();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
ctx.move_resource("x", "y", "line 5").expect("unwrap");
let x = ctx.get_resource("x").expect("unwrap");
assert_eq!(x.ownership, Ownership::Moved);
}
#[test]
fn test_use_after_move_fails() {
let mut ctx = LinearContext::new();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
ctx.move_resource("x", "y", "line 5").expect("unwrap");
let result = ctx.use_resource("x", "line 10");
assert!(matches!(result, Err(LinearError::UseAfterMove { .. })));
}
#[test]
fn test_scope_tracking() {
let mut ctx = LinearContext::new();
ctx.enter_scope();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
ctx.use_resource("x", "line 5").expect("unwrap");
ctx.exit_scope().expect("unwrap");
}
#[test]
fn test_scope_with_unused_linear() {
let mut ctx = LinearContext::new();
ctx.enter_scope();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty, "line 1").expect("unwrap");
let result = ctx.exit_scope();
assert!(result.is_err());
}
#[test]
fn test_statistics() {
let mut ctx = LinearContext::new();
ctx.create_resource("a", LinearType::linear("T"), "1")
.expect("unwrap");
ctx.create_resource("b", LinearType::linear("T"), "2")
.expect("unwrap");
ctx.create_resource("c", LinearType::linear("T"), "3")
.expect("unwrap");
ctx.use_resource("a", "10").expect("unwrap");
ctx.move_resource("b", "d", "20").expect("unwrap");
let stats = ctx.statistics();
assert_eq!(stats.total, 3);
assert_eq!(stats.used, 1);
assert_eq!(stats.unused, 1);
assert_eq!(stats.moved, 1);
}
#[test]
fn test_registry_builtins() {
let registry = LinearTypeRegistry::with_builtins();
assert!(registry.contains("GpuTensor"));
assert!(registry.contains("FileHandle"));
assert!(registry.contains("NetworkConnection"));
let gpu = registry.get("GpuTensor").expect("unwrap");
assert_eq!(gpu.kind, LinearKind::Linear);
}
#[test]
fn test_duplicate_resource() {
let mut ctx = LinearContext::new();
let ty = LinearType::linear("Tensor");
ctx.create_resource("x", ty.clone(), "line 1")
.expect("unwrap");
let result = ctx.create_resource("x", ty, "line 5");
assert!(matches!(result, Err(LinearError::DuplicateResource { .. })));
}
#[test]
fn test_unknown_resource() {
let mut ctx = LinearContext::new();
let result = ctx.use_resource("unknown", "line 1");
assert!(matches!(result, Err(LinearError::UnknownResource { .. })));
}
#[test]
fn test_linear_type_with_tags() {
let ty = LinearType::linear("Resource")
.with_tag("gpu")
.with_tag("memory");
assert_eq!(ty.tags.len(), 2);
assert!(ty.tags.contains(&"gpu".to_string()));
}
}