1 files changed, 159 insertions, 0 deletions
diff --git a/src/Language/Fiddle/Compiler/Backend/Internal/FragTree.hs b/src/Language/Fiddle/Compiler/Backend/Internal/FragTree.hs
new file mode 100644
index 0000000..bc55c5a
--- /dev/null
+++ b/src/Language/Fiddle/Compiler/Backend/Internal/FragTree.hs
@@ -0,0 +1,159 @@
+{-# LANGUAGE DeriveFoldable #-}
+{-# LANGUAGE DeriveTraversable #-}
+
+-- | Module for a 'FragTree' data structure. This data structure represents
+-- "fragments" of a file, which are logical points in a file or buffer.
+--
+-- A fragment represents an indexed position in a file, but the position
+-- is not necessarily discrete. Instead, fragments can be positioned
+-- relative to each other (e.g., above or below other fragments),
+-- allowing for arbitrary resolution within the file.
+module Language.Fiddle.Compiler.Backend.Internal.FragTree
+  ( -- * Types
+    Fragment,
+    FragTree,
+
+    -- * Fragment Constructors
+    center,
+    -- | A fragment at the center of the file.
+    above,
+    -- | A fragment above a given fragment.
+    below,
+    -- | A fragment below a given fragment.
+    -- * FragTree Operations
+    updateWithReturn,
+    -- | Update a fragment in the tree and return a custom value.
+    Language.Fiddle.Compiler.Backend.Internal.FragTree.lookup,
+    -- | Lookup a value by fragment.
+    update,
+    -- | Update or insert a value into the tree.
+    insertOrReplace,
+    -- | Insert a value at a fragment, replacing any existing value.
+    insertAppend,
+    -- | Insert a value by appending (using 'Monoid') if a value already exists.
+    insertWith,
+    -- | Insert a value, combining it with an existing value using a custom function.
+    delete,
+    -- | Delete a value from the tree at a specific fragment.
+    singleton,
+    -- | Create a tree with a single fragment and value.
+  )
+where
+
+-- | Direction of a fragment relative to another fragment.
+data Loc = Above | Below
+
+-- | A 'Fragment' is a position in a file or buffer that can be relative
+-- to other fragments (above or below).
+newtype Fragment = Fragment [Loc]
+
+-- | Fragment that represents the center of the file.
+center :: Fragment
+center = Fragment []
+
+-- | Create a fragment above the given fragment.
+above :: Fragment -> Fragment
+above (Fragment l) = Fragment $ Above : l
+
+-- | Create a fragment below the given fragment.
+below :: Fragment -> Fragment
+below (Fragment l) = Fragment $ Below : l
+
+-- | A 'FragTree' represents a binary tree where each node holds a fragment's value.
+-- Each tree node has two children, representing the fragments above and below
+-- the current fragment. Leaves represent empty nodes.
+data FragTree a
+  = FragTree (FragTree a) (FragTree a) (Maybe a)
+  | FragLeaf
+  deriving (Functor)
+
+instance Foldable FragTree where
+  foldMap _ FragLeaf = mempty
+  foldMap fn (FragTree up down t) =
+    foldMap fn up <> foldMap fn t <> foldMap fn down
+
+instance Traversable FragTree where
+  traverse _ FragLeaf = pure FragLeaf
+  traverse fn (FragTree up down t) = do
+    (\a b c -> FragTree a c b)
+      <$> traverse fn up
+      <*> traverse fn t
+      <*> traverse fn down
+
+-- | 'Semigroup' instance for 'FragTree', merging two trees by combining values
+-- from corresponding fragments, using the 'Semigroup' instance of the contained values.
+instance (Semigroup a) => Semigroup (FragTree a) where
+  (<>) (FragTree ab1 be1 v1) (FragTree ab2 be2 v2) =
+    FragTree (ab1 <> ab2) (be1 <> be2) (v1 <> v2)
+  (<>) FragLeaf f = f
+  (<>) f FragLeaf = f
+
+-- | 'Monoid' instance for 'FragTree'. The empty tree is represented by 'FragLeaf'.
+instance (Semigroup a) => Monoid (FragTree a) where
+  mempty = FragLeaf
+
+-- | 'updateWithReturn' updates a 'FragTree' at the specified 'Fragment',
+-- returning both a custom value and the updated tree.
+-- If the fragment doesn't exist, it is created.
+updateWithReturn :: (Maybe a -> (b, Maybe a)) -> Fragment -> FragTree a -> (b, FragTree a)
+updateWithReturn fn (Fragment (reverse -> loc)) = update' fn loc
+  where
+    -- Helper function for traversing and updating the tree.
+    update' fn [] FragLeaf = clean . FragTree FragLeaf FragLeaf <$> fn Nothing
+    update' fn [] (FragTree a b v) = clean . FragTree a b <$> fn v
+    update' fn (Above : ls) (FragTree a b v) =
+      fmap clean $
+        (\a' -> FragTree a' b v) <$> update' fn ls a
+    update' fn (Below : ls) (FragTree a b v) =
+      fmap clean $
+        (\b' -> FragTree a b' v) <$> update' fn ls b
+    update' fn (Above : ls) FragLeaf =
+      fmap clean $
+        (\a' -> FragTree a' FragLeaf Nothing) <$> update' fn ls FragLeaf
+    update' fn (Below : ls) FragLeaf =
+      fmap clean $
+        (\b' -> FragTree FragLeaf b' Nothing) <$> update' fn ls FragLeaf
+
+-- | Create a 'FragTree' with a single 'Fragment' and associated value.
+singleton :: Fragment -> a -> FragTree a
+singleton f a = insertOrReplace f a FragLeaf
+
+-- | Lookup a value in the 'FragTree' by its 'Fragment'. Returns 'Nothing'
+-- if the fragment does not exist.
+lookup :: Fragment -> FragTree a -> Maybe a
+lookup l = fst . updateWithReturn (\a -> (a, a)) l
+
+-- | Update a 'FragTree' at the given 'Fragment'. If the fragment does not exist,
+-- it is created. The function is applied to the existing value (or 'Nothing' if no value).
+update :: (Maybe a -> Maybe a) -> Fragment -> FragTree a -> FragTree a
+update fn fr = snd . updateWithReturn (\ma -> ((), fn ma)) fr
+
+-- | Insert a value into the 'FragTree', replacing any existing value at the given fragment.
+insertOrReplace :: Fragment -> a -> FragTree a -> FragTree a
+insertOrReplace l a = update (const $ Just a) l
+
+-- | Delete a fragment from the 'FragTree', setting its value to 'Nothing'.
+delete :: Fragment -> FragTree a -> FragTree a
+delete = update (const Nothing)
+
+-- | Insert a value by appending (using the 'Monoid' instance) to an existing value
+-- at the given fragment. If no value exists, the provided value is inserted.
+insertAppend :: (Monoid a) => Fragment -> a -> FragTree a -> FragTree a
+insertAppend l a = update (<> Just a) l
+
+-- | Insert a value using a custom function to combine it with an existing value
+-- at the given fragment. If no val
+insertWith :: (a -> a -> a) -> Fragment -> a -> FragTree a -> FragTree a
+insertWith fn fr v =
+  update
+    ( \case
+        Nothing -> Just v
+        Just a -> Just (fn a v)
+    )
+    fr
+
+-- | Internal function for cleaning up a 'FragTree' by collapsing empty nodes
+-- (i.e., 'FragTree' nodes with no values and no children) into 'FragLeaf'.
+clean :: FragTree a -> FragTree a
+clean (FragTree FragLeaf FragLeaf Nothing) = FragLeaf
+clean x = x