Add support for conditionals.

diff --git a/LTA/Symbolic.hs b/LTA/Symbolic.hs
index b52cfc7bc8038a72b33307c5f22a74b6226dec86..5136006fd7adcf0277978d592c985b90a0114163 100644 (file)
--- a/LTA/Symbolic.hs
+++ b/LTA/Symbolic.hs
@@ -2,18 +2,28 @@
 
 module LTA.Symbolic
   ( Expr(..)
+  , Cond(..)
   , Constant(..)
   , Literal(..)
+  , CompareOp(..)
   , UnaryFunction(..)
+  , buildConditional
   , pow
   , simplify
+  , (~==)
+  , (~<)
+  , (~<=)
+  , (~>)
+  , (~>=)
   ) where
 
 import Control.Applicative ((<$>))
 import Data.Ratio (numerator, denominator)
 import Data.List (foldl')
 import Data.Map (Map)
+import Data.Set (Set)
 import qualified Data.Map as Map
+import qualified Data.Set as Set
 
 data SumTag
 data ProductTag
@@ -86,6 +96,24 @@ data Expr
   | UnaryFunction UnaryFunction Expr
   | Div Expr Expr
   | Mod Expr Expr
+  | Conditional Cond Expr Expr
+  deriving (Eq, Ord, Show)
+
+data CompareOp
+ = Equal
+ | LessThan
+ | LessThanEqual
+ | GreaterThan
+ | GreaterThanEqual
+ deriving (Eq, Ord, Show)
+
+data Cond
+  = TrueC
+  | FalseC
+  | Compare CompareOp Expr Expr
+  | And Cond Cond
+  | Or Cond Cond
+  | Not Cond
   deriving (Eq, Ord, Show)
 
 data Literal
@@ -108,6 +136,26 @@ data UnaryFunction
   | Signum
   deriving (Eq, Ord, Show)
 
+buildConditional :: Expr -> [(Cond, Expr)] -> Expr
+buildConditional base [] = base
+buildConditional base ((cond, expr) : cases) =
+  Conditional cond expr (buildConditional base cases)
+
+(~<) :: Expr -> Expr -> Cond
+(~<) = Compare LessThan
+
+(~<=) :: Expr -> Expr -> Cond
+(~<=) = Compare LessThanEqual
+
+(~==) :: Expr -> Expr -> Cond
+(~==) = Compare Equal
+
+(~>=) :: Expr -> Expr -> Cond
+(~>=) = Compare GreaterThanEqual
+
+(~>) :: Expr -> Expr -> Cond
+(~>) = Compare GreaterThan
+
 simplify :: Expr -> Expr
 simplify (Sum pairSeq) = rebuild $ normalise pairSeq
 simplify (Product pairSeq) = rebuild $ normalise pairSeq

diff --git a/src/Main.hs b/src/Main.hs
index bb5d7f9065ba36a929c94cd677ab9d6de42bff1a..952e49dcb57314bb8c33b4bb494a09b54e5566f5 100644 (file)
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -3,7 +3,7 @@ module Main (main) where
 import LTA.Symbolic
 
 main :: IO()
-main = putStrLn $ show dft
+main = putStrLn . show $ dft (IntegerSymbol "N")
 
 row :: Expr
 row = IntegerSymbol "row"
@@ -11,9 +11,6 @@ row = IntegerSymbol "row"
 col :: Expr
 col = IntegerSymbol "col"
 
-n :: Expr
-n = IntegerSymbol "N"
-
 e :: Expr
 e = Constant Euler
 
@@ -23,5 +20,29 @@ i = Constant ImaginaryUnit
 pi :: Expr
 pi = Constant Pi
 
-dft :: Expr
-dft = pow e ((- i * 2 * Main.pi * row * col) / n)
+dft :: Expr -> Expr
+dft size = pow e ((- i * 2 * Main.pi * row * col) / size)
+
+isEven :: Expr -> Cond
+isEven n = (n `Mod` 2) ~== 0
+
+pad :: Expr -> Expr
+pad size = Conditional (isEven size) (evenPad size) (oddPad size)
+
+evenPad :: Expr -> Expr
+evenPad size = buildConditional 0 conditions where
+  split = size `Div` 2
+  conditions = [ ((row ~== col) `And` (row ~< split), 1)
+               , ((row ~== col) `And` (row ~== split), 0.5)
+               , ((row ~> split) `And` (row ~< (size + split)), 0)
+               , ((row ~== (size + split)) `And` (col ~== split), 0.5)
+               , ((row ~> (size + split)) `And` (row ~== (col + size + split)), 1)
+               ]
+
+oddPad :: Expr -> Expr
+oddPad size = buildConditional 0 conditions where
+  split = (size + 1) `Div` 2
+  conditions = [ ((row ~== col) `And` (row ~< split), 1)
+               , ((row ~>= split) `And` (row ~< (size + split)), 0)
+               , ((row ~>= (size + split)) `And` (row ~== (col + size + split)), 1)
+               ]