import Mathlib.Data.Finset.Basic

/-!

Base language.

This file defines the core object language for the program analysis and
transformation. It's a very basic imperative language. The `Spa/Language/Tagged/Basic.lean`
file provides an auto-derived version of the `Expr`, `BasicStmt`, and `Stmt` data
types with unique IDs per condtructor, enabling in-AST pointers.

-/

namespace Spa

/-- A value-producing expression. Currently, this cannot have side effects. -/
inductive Expr where
  | add (e₁ e₂ : Expr)
  | sub (e₁ e₂ : Expr)
  | var (x : String)
  | num (n : ℕ)
  deriving DecidableEq

/-- A statement that cannot alter control flow (and thus, can be part of a basic block).

    This differs from, e.g., a loop, which can cause execution to jump to its top several times. -/
inductive BasicStmt where
  | assign (x : String) (e : Expr)
  | noop
  deriving DecidableEq

/-- Any statements, which may or may not change program state (variable assignments). -/
inductive Stmt where
  | basic (bs : BasicStmt)
  | andThen (s₁ s₂ : Stmt)
  | ifElse (e : Expr) (s₁ s₂ : Stmt)
  | whileLoop (e : Expr) (s : Stmt)
  deriving DecidableEq

/-- Variables mentioned in this expression. -/
def Expr.vars : Expr → Finset String
  | .add l r => l.vars ∪ r.vars
  | .sub l r => l.vars ∪ r.vars
  | .var s => {s}
  | .num _ => ∅

/-- Variables assigned or mentioned in this basic statement. -/
def BasicStmt.vars : BasicStmt → Finset String
  | .assign x e => {x} ∪ e.vars
  | .noop => ∅

/-- Variables assigned or mentioned in this statement. -/
def Stmt.vars : Stmt → Finset String
  | .basic bs => bs.vars
  | .andThen s₁ s₂ => s₁.vars ∪ s₂.vars
  | .ifElse e s₁ s₂ => (e.vars ∪ s₁.vars) ∪ s₂.vars
  | .whileLoop e s => e.vars ∪ s.vars

end Spa