Rotate.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}

module Zwirn.Language.Rotate
  ( runRotate,
    runRotateUnsafe,
    RotationError (..),
  )
where

{-
    Rotate.hs - syntax tree rotation, code adapted from
    https://gist.github.com/heitor-lassarote/b20d6da0a9042d31e439befb8c236a4e
    Copyright (C) 2023, Martin Gius

    This library is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This library is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this library.  If not, see <http://www.gnu.org/licenses/>.
-}

import Control.Monad.Except
import Control.Monad.Identity
import Zwirn.Language.Location
import Zwirn.Language.Simple
import Zwirn.Language.Syntax

ops :: [Declaration]
ops =
  [ ("(*)", Fixity LeftA 9),
    ("(/)", Fixity LeftA 9),
    ("(.)", Fixity RightA 9),
    ("(+)", Fixity LeftA 6),
    ("(-)", Fixity LeftA 6),
    ("($)", Fixity RightA 0),
    ("($:)", Fixity RightA 0),
    ("(&:)", Fixity RightA 0),
    ("($!)", Fixity RightA 0),
    ("(#!)", Fixity RightA 0),
    ("($|)", Fixity RightA 0),
    ("(|$)", Fixity RightA 0),
    ("(#)", Fixity RightA 3),
    ("(&)", Fixity RightA 2),
    ("(++)", Fixity RightA 5),
    -- arithmetic
    ("(|-)", Fixity LeftA 6),
    ("(-|)", Fixity LeftA 6),
    ("(|+)", Fixity LeftA 6),
    ("(+|)", Fixity LeftA 6),
    ("(|*)", Fixity LeftA 7),
    ("(*|)", Fixity LeftA 7),
    ("(//)", Fixity LeftA 7),
    ("(|/)", Fixity LeftA 7),
    ("(/|)", Fixity LeftA 7),
    -- ord
    ("(<=)", Fixity NonA 4),
    ("(>=)", Fixity NonA 4),
    ("(&&)", Fixity RightA 3),
    ("(||)", Fixity RightA 3),
    ("(==)", Fixity NonA 4)
  ]

defaultFixity :: Fixity
defaultFixity = Fixity LeftA 8

newtype RotationError = RotationError SrcLoc

instance Show RotationError where
  show (RotationError (SrcLoc p)) = "Could not resolve operator precedence at location " ++ show p
  show (RotationError NoLoc) = "Could not resolve operator precedence"

type Rotate a = ExceptT RotationError Identity a

-- | Describes which action the rotation algorithm should use.
data Rotation
  = -- | Fail due to the mixing of incompatible operators.
    Fail
  | -- | Keep the tree as it is.
    Keep
  | -- | Balance the tree to the left.
    Rotate

runRotate :: LocSimpleTerm -> Either RotationError LocSimpleTerm
runRotate t = runIdentity $ runExceptT $ rotate t

runRotateUnsafe :: LocSimpleTerm -> LocSimpleTerm
runRotateUnsafe t = case runRotate t of
  Left err -> error $ show err
  Right r -> r

-- | The Happy parser is written in a way so that it will always create a right-balanced AST.
-- We compare the operators and indicate how to rotate the tree.
shouldRotate :: Fixity -> Fixity -> Rotation
shouldRotate (Fixity a p) (Fixity a' p') = case compare p p' of
  LT -> Keep
  EQ -> case (a, a') of
    (LeftA, LeftA) -> Rotate
    (RightA, RightA) -> Keep
    (_, _) -> Fail
  GT -> Rotate

-- | Rebalances the tree to respect the associativity and precedence of the
-- parsed operators.

-- Not very efficient, but enough for demonstration purposes.
findOp :: LocVar -> Rotate Fixity
findOp (Located _ o) = case lookup o ops of
  Just d -> return d
  Nothing -> return defaultFixity

rotate :: LocSimpleTerm -> Rotate LocSimpleTerm
rotate (Located p1 (SInfix l op r)) = do
  -- Rotating the left side is unneeded since this grammar is very simple.
  -- This is because trees are always right-balanced and the left side is
  -- always an atom.
  lRotated <- rotate l
  rRotated <- rotate r
  case rRotated of
    (Located p2 (SInfix l' op' r')) -> do
      opDec <- findOp op
      opDec' <- findOp op'
      case shouldRotate opDec opDec' of
        Fail -> throwError (RotationError $ lLoc op)
        Keep -> return $ Located p1 $ SInfix lRotated op rRotated
        Rotate -> return $ Located p2 $ SInfix (Located p1 $ SInfix lRotated op l') op' r'
    _ -> return $ Located p1 $ SInfix lRotated op rRotated
rotate (Located p (SApp l r)) = do
  lRotated <- rotate l
  rRotated <- rotate r
  return $ Located p $ SApp lRotated rRotated
rotate (Located p e@(SVar _)) = return $ Located p e
rotate (Located p e@(SText _)) = return $ Located p e
rotate (Located p e@(SNum _)) = return $ Located p e
rotate (Located p SRest) = return $ Located p SRest
rotate (Located p (SSeq ts)) = fmap (Located p . SSeq) (mapM rotate ts)
rotate (Located p (SStack ts)) = fmap (Located p . SStack) (mapM rotate ts)
rotate (Located p (SChoice n ts)) = fmap (Located p . SChoice n) (mapM rotate ts)
rotate (Located p (SLambda vs t)) = Located p . SLambda vs <$> rotate t
rotate (Located p (SBracket t)) = fmap (Located p . SBracket) (rotate t)