theory Scratch
imports "~~/src/HOL/Library/FSet"
begin

datatype val = VNat nat | VRel "(val \<times> val) fset"

lemma sum_nat_le_single:
  fixes x :: nat
  assumes y: "y \<in> S" and x: "x \<le> f y" and finite: "finite S"
  shows "x \<le> sum f S"
proof -
  have "sum f S = sum f (insert y S)"
    using y by (metis insert_absorb)
  with finite have "sum f S = f y + sum f (S - {y})"
    by (metis sum.insert_remove)
  with x show ?thesis
    by linarith
qed

context includes fset.lifting begin

lift_definition fsum :: "('a \<Rightarrow> 'b) \<Rightarrow> 'a fset \<Rightarrow> 'b::comm_monoid_add" is sum .

lemma fsum_cong[fundef_cong]: "A = B \<Longrightarrow> (\<And>x. x |\<in>| B \<Longrightarrow> g x = h x) \<Longrightarrow> fsum g A = fsum h B"
by transfer' (rule sum.cong)

function (sequential) count :: "val \<Rightarrow> nat" where
"count (VNat n) = 1" |
"count (VRel t) = fsum (\<lambda>(v, u). count v + count u) t"
  by pat_completeness auto

termination
  apply (relation "measure size")
  subgoal by simp
  subgoal
    apply simp
    apply transfer
    apply (rule le_imp_less_Suc)
    apply (rule sum_nat_le_single)
    by auto
  subgoal
    apply simp
    apply transfer
    apply (rule le_imp_less_Suc)
    apply (rule sum_nat_le_single)
    by auto
  done

end

primrec count' :: "val \<Rightarrow> nat" where
"count' (VNat n) = 1" |
"count' (VRel t) = fsum (\<lambda>(v, u). v + u) (map_prod count' count' |`| t)"

end