Code review: Defining a partial function, then proving total

Stream: new members

[C7X]

Hi!

I have a minimal working example here where I define a partial function expand', prove its totality, then use that to define a total version expand of expand'. From then on I will only use expand without referencing expand'. Is this bad practice?

import Mathlib

def PSeq : Type := List ℕ

def bad_root (s : PSeq) : Option (Fin (s.length)) :=
  match s with
  | [] => none
  | x :: s' =>
    (x :: s').findFinIdx? (· < x)

def good_part (s : PSeq) : Option PSeq :=
  match bad_root s with
  | some br => some (s.take br)
  | none => none

-- First define expansion, then prove it total

def expand' (s : PSeq) : Option PSeq :=
  if (bad_root s).isSome
    then good_part s
    else s

lemma expand'_total (s : PSeq) : (expand' s).isSome := by {
  simp [expand']
  split
  simp [good_part]
  split
  · trivial
  rename (bad_root s).isSome = true => h2
  rename bad_root s = none => h3
  rw [← h2]
  rw [h3]
  trivial
  trivial
}

def expand (s : PSeq) : PSeq :=
  (expand' s).get (expand'_total s)

-- From here on, expand' will never be used, only expand will.

[Kenny Lau]

as long as you have good API, you can do anything you want

[Eric Wieser]

Though in this case you enumerate the list twice rather than once, which matters if you #eval