prog_synthesis/02.hs

import Data.List (elemIndex, sort, find)
import Data.Maybe (isJust)

infixl 4 :+:

data IntExpr = Zero
             | Succ IntExpr
             | Len ListExpr
             | FirstZero ListExpr
    deriving (Read, Show, Eq)

data ListExpr = Sort ListExpr
              | SubList ListExpr IntExpr IntExpr
              | ListExpr :+: ListExpr
              | Recursive ListExpr
              | ZeroList
              | InList
    deriving (Read, Show, Eq)

data Env = Env {input :: [Int], prog :: ListExpr, steps :: Int}

stepInEnv :: Env -> Env
stepInEnv env@(Env {input, prog, steps}) = env { steps = steps - 1 }

execInt :: Env -> IntExpr -> Maybe Int
execInt (Env {input, prog, steps=0}) _ = Nothing
execInt _ Zero = Just 0
execInt env (Succ expr) = (+) 1 <$> execInt (stepInEnv env) expr
execInt env (Len listExpr) = Just $ length $ execList (stepInEnv env) listExpr
execInt env (FirstZero listExpr) = do value <- execList (stepInEnv env) listExpr
                                      0 `elemIndex` value

-- TODO: limit execution steps
execList :: Env -> ListExpr -> Maybe [Int]
execList (Env {input, prog, steps=0}) _ = Nothing
execList env (Sort listExpr) = sort <$> execList (stepInEnv env) listExpr
execList env (SubList listExpr exprFrom exprTo) = do valFrom  <- execInt (stepInEnv env) exprFrom
                                                     valTo <- execInt (stepInEnv env) exprTo
                                                     listValue <- execList (stepInEnv env) listExpr
                                                     return $ drop valFrom $ take valTo listValue
execList env (exprLeft :+: exprRight) = do valLeft <- execList (stepInEnv env) exprLeft
                                           valRight <- execList (stepInEnv env) exprRight
                                           return $ valLeft ++ valRight
execList env (Recursive listExpr) = Nothing -- do listValue <- execList (stepInEnv env) listExpr
                                    --   if null listValue then Just [] else execList (stepInEnv $ env {input = listValue}) (prog env)
execList _ ZeroList = Just [0]
execList (Env {input, prog, steps}) InList = Just input

execProg :: [Int] -> ListExpr -> Maybe [Int]
execProg input expr = execList (Env {input, prog=expr, steps=100}) expr

data AllExprs = AllExprs {ints :: [IntExpr], lists :: [ListExpr]}

terminals :: AllExprs
terminals = AllExprs {ints = [Zero], lists = [ZeroList, InList]}

nextExprs :: AllExprs -> AllExprs
nextExprs (AllExprs {ints, lists}) = AllExprs {ints = map Succ ints ++ map Len lists ++ map FirstZero lists ++ ints,
                                               lists = map Sort lists ++ zipWith3 SubList lists ints ints ++ zipWith (:+:) lists lists ++ map Recursive lists ++ lists}

data Example = Example {exampleInput :: [Int], exampleOutput :: [Int]}

-- check expr on all examples
isCorrect :: [Example] -> ListExpr -> Bool
isCorrect examples expr = all (\Example {exampleInput, exampleOutput} -> execProg exampleInput expr == Just exampleOutput) examples

isValid :: [Example] -> ListExpr -> Bool
isValid examples expr = all (\Example {exampleInput, exampleOutput} -> isJust $ execProg exampleInput expr) examples

-- check are exprs produce same results on all the examples
areSame :: [Example] -> ListExpr -> ListExpr -> Bool
areSame examples exprLeft exprRight = all (\Example {exampleInput, exampleOutput} -> execProg exampleInput exprLeft == execProg exampleInput exprRight ) examples

upSyntesisStep :: [Example] -> AllExprs -> Either AllExprs ListExpr
upSyntesisStep examples allExprs@(AllExprs {ints, lists}) = let lists = filter (isValid examples) lists in -- exclude invalid fragments
                                                            -- let lists = foldl (\acc expr -> if any (areSame examples expr) acc then acc else expr : acc) [] lists' in -- merge same values
                                                            case find (isCorrect examples) lists of
                                                              Just answer -> Right answer
                                                              Nothing -> Left $ nextExprs allExprs

upSyntesisRec :: [Example] -> Int -> AllExprs  -> Maybe ListExpr
upSyntesisRec _ 0 _ = Nothing
upSyntesisRec examples steps exprs = case upSyntesisStep examples exprs of
                                          Right answer -> Just answer
                                          Left exprs' -> upSyntesisRec examples (steps - 1) exprs'

upSyntesis :: [Example] -> Maybe ListExpr
upSyntesis examples = upSyntesisRec examples 1 terminals

-----

syntesisTest1 = upSyntesis [Example {exampleInput = [1,2,3,4,5,6,7,8], exampleOutput = [1,2,3,4,5,6,7,8]}]
syntesisTest2 = upSyntesis [Example {exampleInput = [1,2,3,4,5,6,7,8], exampleOutput = [8,7,6,5,4,3,2,1]}]